Code agent

import json
import os
import subprocess
import sys

from openai import OpenAI

client = OpenAI(base_url="https://api.hcompany.ai/v1/", api_key=os.environ["HAI_API_KEY"])
MODEL = "holo4-35b-a3b"


def fn(name: str, description: str, **properties) -> dict:
    parameters = {"type": "object", "properties": properties, "required": list(properties)}
    return {"type": "function", "function": {"name": name, "description": description, "parameters": parameters}}


tools = [
    fn("run_python", "Run a Python script and return its stdout and stderr. Print anything you want to see.", code={"type": "string", "description": "Python source code"}),
    fn("answer", "Provide a final answer", content={"type": "string", "description": "The answer content"}),
]


def run_python(code: str) -> str:
    try:
        proc = subprocess.run([sys.executable, "-c", code], capture_output=True, text=True, timeout=30)
    except subprocess.TimeoutExpired:
        return "Timed out after 30 seconds"
    return (proc.stdout + proc.stderr)[-4000:] or "(no output)"


def run(task: str, max_steps: int = 10) -> str:
    messages = [
        {"role": "system", "content": "You solve tasks by writing and running Python. Verify results by running code before answering."},
        {"role": "user", "content": task},
    ]
    for _ in range(max_steps):
        resp = client.chat.completions.create(
            model=MODEL,
            messages=messages,
            tools=tools,
            tool_choice="required",
            temperature=0.8,
            extra_body={"chat_template_kwargs": {"enable_thinking": True}},
        )
        msg = resp.choices[0].message
        messages.append({"role": "assistant", "content": msg.content, "tool_calls": msg.tool_calls})
        if not msg.tool_calls:
            continue
        call = msg.tool_calls[0]
        args = json.loads(call.function.arguments)
        if call.function.name == "answer":
            return args["content"]
        print(f"--- run_python\n{args['code']}")
        result = run_python(args["code"])
        print(f"--- output\n{result}")
        messages.append({"role": "tool", "tool_call_id": call.id, "content": result})
    return "Step budget exhausted"


if __name__ == "__main__":
    print(run("How many prime numbers are there below 1,000,000, and what is the largest one?"))

--- run_python
def count_primes_and_find_largest(limit):
    primes = []
    for num in range(2, limit):
        is_prime = True
        for prime in primes:
            if prime * prime > num:
                break
            if num % prime == 0:
                is_prime = False
                break
        if is_prime:
            primes.append(num)
    return len(primes), primes[-1]

count, largest = count_primes_and_find_largest(1000000)
print(f"Count: {count}, Largest: {largest}")
--- output
Count: 78498, Largest: 999983

There are 78,498 prime numbers below 1,000,000, and the largest one is 999,983.