Coding Agent: Mercury + Tool Calling

"""
Mercury 2 tool-calling example.

Usage:
    export INCEPTION_API_KEY=...
    python coding-agent.py
"""

import json
import os
import requests


API_URL = "https://api.inceptionlabs.ai/v1/chat/completions"
MODEL = "mercury-2"
HEADERS = {
    "Authorization": f"Bearer {os.environ['INCEPTION_API_KEY']}",
    "Content-Type": "application/json",
}


# ---------------------------------------------------------------------------
# Mock filesystem the "tools" will pretend to read.
# ---------------------------------------------------------------------------
FAKE_FS = {
    "/repo": ["main.py", "stats.py"],
    "/repo/stats.py": (
        "def compute_average(values):\n"
        "    \"\"\"Return the arithmetic mean of a list of numbers.\"\"\"\n"
        "    if not values:\n"
        "        return 0.0\n"
        "    return sum(values) / len(values)\n"
    ),
    "/repo/main.py": (
        "from stats import compute_average\n"
        "\n"
        "print(compute_average([1, 2, 3]))\n"
    ),
}


def run_list_directory(path: str) -> str:
    entries = FAKE_FS.get(path)
    if entries is None or not isinstance(entries, list):
        return json.dumps({"error": f"not a directory: {path}"})
    return json.dumps({"path": path, "entries": entries})


def run_read_file(path: str) -> str:
    contents = FAKE_FS.get(path)
    if contents is None or isinstance(contents, list):
        return json.dumps({"error": f"not a file: {path}"})
    return json.dumps({"path": path, "contents": contents})


TOOL_RUNNERS = {
    "list_directory": lambda args: run_list_directory(args["path"]),
    "read_file": lambda args: run_read_file(args["path"]),
}


# ---------------------------------------------------------------------------
# Tool schemas — the OpenAI tools[] format (which Mercury 2 also accepts).
# ---------------------------------------------------------------------------
TOOLS = [
    {
        "type": "function",
        "function": {
            "name": "list_directory",
            "description": "List the entries (files and subdirectories) inside a directory.",
            "parameters": {
                "type": "object",
                "properties": {
                    "path": {
                        "type": "string",
                        "description": "Absolute directory path, e.g. '/repo'.",
                    },
                },
                "required": ["path"],
                "additionalProperties": False,
            },
        },
    },
    {
        "type": "function",
        "function": {
            "name": "read_file",
            "description": "Return the full text contents of a file.",
            "parameters": {
                "type": "object",
                "properties": {
                    "path": {
                        "type": "string",
                        "description": "Absolute file path, e.g. '/repo/stats.py'.",
                    },
                },
                "required": ["path"],
                "additionalProperties": False,
            },
        },
    },
]


SYSTEM_PROMPT = (
    "You are a coding assistant with access to a small filesystem via two tools: "
    "`list_directory` and `read_file`. Use tools to explore the repo before answering. "
    "When you have enough information, reply in plain text with the answer."
)

USER_PROMPT = (
    "The repo lives at `/repo`. Find the definition of the function "
    "`compute_average` and quote its body back to me."
)


def chat_completion(messages):
    """POST to the Inception chat/completions endpoint and return the JSON body."""
    payload = {
        "model": MODEL,
        "messages": messages,
        "tools": TOOLS,
    }
    resp = requests.post(API_URL, headers=HEADERS, json=payload)
    resp.raise_for_status()
    return resp.json()


def main():
    messages = [
        {"role": "system", "content": SYSTEM_PROMPT},
        {"role": "user", "content": USER_PROMPT},
    ]

    for turn in range(3):
        print(f"\n========== turn {turn} ==========")
        body = chat_completion(messages)
        choice = body["choices"][0]
        msg = choice["message"]
        finish = choice.get("finish_reason")
        tool_calls = msg.get("tool_calls") or []
        print(f"finish_reason: {finish}")

        if tool_calls:
            # add tool calls to messages history
            messages.append({
                "role": "assistant",
                "content": msg.get("content"),
                "tool_calls": tool_calls,
            })

            for tc in tool_calls:
                print(f"tc: {tc}")
                fn = tc["function"]
                print(f"fn: {fn}")
                name = fn["name"]
                raw_args = fn["arguments"]
                print(f"  tool_call id={tc['id']}")
                print(f"    name: {name}")
                print(f"    args: {raw_args}")

                try:
                    parsed_args = json.loads(raw_args)
                except json.JSONDecodeError:
                    parsed_args = {}

                runner = TOOL_RUNNERS.get(name)
                if runner is None:
                    result = json.dumps({"error": f"unknown tool {name}"})
                else:
                    result = runner(parsed_args)
                print(f"    result: {result}")
                # add tool result to messages history
                messages.append({
                    "role": "tool",
                    "tool_call_id": tc["id"],
                    "content": result,
                })
        else:
            # answer — loop is done.
            print(f"assistant output:\n{msg.get('content')}")
            break


if __name__ == "__main__":
    main()