LLM 代码练习——调用 API 进行简单聊天

一个简单的调用 LLM API 的大模型单次对话脚本，通过 --input 和 --prompt 参数提供输入文件和系统提示文件，输出到标准输出流。

脚本调用 LLM 的方式不依赖 openai 等提供的第三方库，完全使用 Python 标准库实现。

chat-with-prompt.py

#!/usr/bin/env python3

import json
import os
import sys
import re
import argparse
import urllib.request
import urllib.error
import logging

logger = logging.getLogger(__name__)


class Colors:
    ENABLED = sys.stdout.isatty()

    RED = "\033[31m"
    GREEN = "\033[32m"
    YELLOW = "\033[33m"
    BLUE = "\033[34m"
    MAGENTA = "\033[35m"
    CYAN = "\033[36m"
    RED_BG = "\033[41m"
    GRAY = "\033[90m"
    RESET = "\033[0m"

    BOLD = "\033[1m"
    DIM = "\033[2m"

    @staticmethod
    def wrap(text: str, color: str) -> str:
        if Colors.ENABLED and color:
            return f"{color}{text}{Colors.RESET}"
        return text


def separator():
    try:
        columns = os.get_terminal_size().columns
    except OSError:
        columns = 80
    return Colors.wrap(f"{'-' * min(columns, 80)}", Colors.DIM)


def render_markdown(text):
    """Simple regex-based markdown renderer for terminal output."""
    if not Colors.ENABLED:
        return text
    # **bold** -> BOLD
    text = re.sub(r"\*\*(.+?)\*\*", f"{Colors.BOLD}\\1{Colors.RESET}", text)
    # *italic* -> DIM
    text = re.sub(r"\*(.+?)\*", f"{Colors.DIM}\\1{Colors.RESET}", text)
    # `code` -> CYAN
    text = re.sub(r"`(.+?)`", f"{Colors.CYAN}\\1{Colors.RESET}", text)
    return text


def read_file(path: str) -> str:
    try:
        with open(path, "r", encoding="utf-8") as f:
            return f.read().strip()
    except Exception as e:
        print(Colors.wrap(f"Error reading file '{path}': {e}", Colors.RED))
        sys.exit(1)


def call_api(
    *,
    headers: dict[str, str],
    base_url: str,
    model: str,
    system_msg: str,
    user_msg: str,
    temperature: float,
    max_tokens: int,
) -> str:
    """Performs the API request via urllib."""
    is_anthropic = "claude" in model.lower()

    if is_anthropic:
        payload = {
            "model": model,
            "system": system_msg,
            "messages": [
                {"role": "user", "content": user_msg},
            ],
            "temperature": temperature,
            "max_tokens": max_tokens,
        }
    else:
        payload = {
            "model": model,
            "messages": [
                {"role": "system", "content": system_msg},
                {"role": "user", "content": user_msg},
            ],
            "temperature": temperature,
            "max_tokens": max_tokens,
        }

    logger.debug(f"API Request to {base_url}")
    logger.debug(f"Payload: {json.dumps(payload, indent=2, ensure_ascii=False)}")
    logger.debug(f"Headers: {json.dumps(headers, indent=2, ensure_ascii=False)}")

    req = urllib.request.Request(
        base_url, data=json.dumps(payload).encode(), headers=headers
    )

    try:
        with urllib.request.urlopen(req) as response:
            res_body = response.read().decode("utf-8")
            logger.debug(f"API Response Raw: {res_body}")

            res_data = json.loads(res_body)

            content = ""
            if "choices" in res_data:
                # OpenAI / Shared Standard Schema
                content = res_data["choices"][0]["message"]["content"]
            elif "content" in res_data:
                # Anthropic Native Schema
                # Extract text from content block list
                content = "".join(
                    [
                        block["text"]
                        for block in res_data["content"]
                        if block.get("type") == "text"
                    ]
                )
            else:
                raise KeyError(f"Unexpected response format: {list(res_data.keys())}")

            # Token Stats Logging
            usage = res_data.get("usage", {})
            p_tokens = usage.get("prompt_tokens") or usage.get("input_tokens", 0)
            c_tokens = usage.get("completion_tokens") or usage.get("output_tokens", 0)

            stats = f"  (Tokens: {p_tokens}↑ {c_tokens}↓ | Temperature: {temperature})"
            print(Colors.wrap(stats, Colors.DIM))

            return content
    except urllib.error.HTTPError as e:
        error_body = e.read().decode()
        raise Exception(f"HTTP {e.code}: {error_body}")
    except Exception as e:
        raise Exception(f"Connection Error: {e}")


def main():
    parser = argparse.ArgumentParser(description="LLM Batch Processor")
    parser.add_argument("--prompt", required=True, help="Path to system prompt file")
    parser.add_argument("--input", required=True, help="Path to user input file")
    parser.add_argument(
        "--output", help="Path to save output. If omitted, prints to stdout."
    )
    parser.add_argument(
        "--model",
        default=os.getenv("MODEL", "openrouter/free"),
        help="Model name (Env: MODEL)",
    )

    parser.add_argument(
        "--base-url",
        default=os.getenv("BASE_URL", "https://openrouter.ai/api/v1"),
        help="API Base URL (Env: BASE_URL)",
    )
    parser.add_argument(
        "--temperature",
        type=float,
        default=0.2,
        help="Sampling temperature (default: 0.2)",
    )
    parser.add_argument(
        "--debug-api", action="store_true", help="Enable verbose API logging to stderr"
    )
    args = parser.parse_args()

    logging.basicConfig(
        level=logging.DEBUG if args.debug_api else logging.WARNING,
        format="[%(levelname)s] %(message)s",
    )

    model = args.model
    base_url = args.base_url.rstrip("/")
    is_anthropic = "claude" in model.lower()
    if is_anthropic:
        base_url = f"{base_url}/messages"
    else:
        base_url = f"{base_url}/chat/completions"

    print(
        f"{Colors.wrap('LLM Batch Processor', Colors.BOLD)} | {Colors.wrap(args.model, Colors.CYAN)}"
    )
    print(separator())

    api_key = os.getenv("OPENROUTER_API_KEY")
    if not api_key:
        print(
            Colors.wrap(
                "⏺ Error: API Key not found in environment variables.", Colors.RED
            )
        )
        sys.exit(1)

    headers = {
        "Content-Type": "application/json",
        "Authorization": f"Bearer {api_key}",
    }
    if is_anthropic:
        headers["anthropic-version"] = "2023-06-01"

    DEFAULT_MAX_TOKENS = 8192

    try:
        system_msg = read_file(args.prompt)
        user_msg = read_file(args.input)

        print(Colors.wrap("⏺ Running API Request...", Colors.YELLOW))

        result = call_api(
            headers=headers,
            base_url=base_url,
            model=model,
            system_msg=system_msg,
            user_msg=user_msg,
            temperature=args.temperature,
            max_tokens=DEFAULT_MAX_TOKENS,
        )

        if args.output:
            with open(args.output, "w", encoding="utf-8", newline="\n") as f:
                f.write(result)
            print(Colors.wrap("⏺ Success!", Colors.GREEN))
            print(Colors.wrap(f"⏺ Result saved to: {args.output}", Colors.DIM))
        else:
            print(Colors.wrap("⏺ Success!", Colors.GREEN))
            print(separator())
            print(render_markdown(result))
            print(separator())

    except Exception as e:
        print(Colors.wrap(f"⏺ Fatal Error: {e}", Colors.RED))
        sys.exit(1)


if __name__ == "__main__":
    main()

输入指令文件

input.txt

I think AI is very important for doctors. It can look at many medical images
and find diseases very fast. This helps doctors save time and makes
fewer mistakes. We did an experiment and the results were good.

prompt 文件

prompt.txt

You are a professional academic editor specialized in computer science.
Your task is to polish the user's paragraph to make it sound more formal,
concise, and suitable for publication in a top-tier journal like Nature or IEEE.
Please provide the polished version and a brief explanation of changes.