from openai import OpenAI

client = OpenAI(
    base_url="http://localhost:8000/v1",
    api_key="sk-local",
)

# 非流式
response = client.chat.completions.create(
    model="rwkv-7",
    messages=[{"role": "user", "content": "你好！"}],
    temperature=0.7,
    max_tokens=500,
)
print(response.choices[0].message.content)

# 流式
stream = client.chat.completions.create(
    model="rwkv-7",
    messages=[{"role": "user", "content": "讲个笑话"}],
    stream=True,
)
for chunk in stream:
    if chunk.choices[0].delta.content:
        print(chunk.choices[0].delta.content, end="", flush=True)

response = client.chat.completions.create(
    model="rwkv-7",
    messages=[
        {"role": "system", "content": "你是一个专业的技术翻译，请将英文翻译为中文。"},
        {"role": "user", "content": "LLMs use multi-head self-attention to process tokens in parallel."},
    ],
)
# 实际 prompt：
# System: 你是一个专业的技术翻译...
#
# User: LLMs use multi-head self-attention...
#
# Assistant:  thinking
#  response

from server import RWKV070ModelLoader, DynamicScheduler
from server.task.manager import get_task_manager

pipeline = RWKV070ModelLoader("model.pth")
scheduler = DynamicScheduler(pipeline)
scheduler.start_daemon()
task_manager = get_task_manager()

# 创建翻译模板：max_tokens=0，只 prefill 不生成
template_task = scheduler.new_task(
    prompt="User: 你是一个专业的中英技术翻译，请翻译以下内容。\n\n",
    max_tokens=0,  # 关键：只 prefill
    temperature=0.3,
)
scheduler.run()  # prefill 完成，state 已就绪
task_manager.put_task("_tech_translator", template_task)

# 每次使用时 fork，state 自动继承，零 prefill 开销
new_id = task_manager.fork_template("_tech_translator")
task = task_manager.get_task(new_id)
task.continue_gen(
    "请翻译：The transformer architecture revolutionized NLP.",
    max_tokens=500,
)
scheduler.add_task(task)

roles = {
    "_coder":       "User: 你是一个资深的 Python 后端工程师。\n\nAssistant: 好的。\n\n",
    "_translator":  "User: 你是一个专业的中英技术翻译。\n\nAssistant: 明白。\n\n",
    "_writer":      "User: 你是一个创意故事作家。\n\nAssistant: 好的，让我们开始。\n\n",
}

for role_id, prompt in roles.items():
    task = scheduler.new_task(prompt=prompt, max_tokens=0)
    scheduler.run()
    task_manager.put_task(role_id, task)

# 随时 fork 使用
coding_task_id = task_manager.fork_template("_coder")

from server import RWKV070ModelLoader, DynamicScheduler
from server.config import settings

# 1. 加载模型
loader = RWKV070ModelLoader(settings.model_path, settings.vocab_path)

# 2. 创建调度器
scheduler = DynamicScheduler(
    loader,
    max_batch_size=256,   # 最大并发
    buffer_size=32,       # 每个 pulse 生成的 token 数
)

# 3. 提交任务
tasks = [
    scheduler.new_task(
        f"User: {prompt}\n\nAssistant:",
        max_tokens=200,
        temperature=0.7,
        top_p=0.3,
        top_k=50,
    )
    for prompt in ["你好，介绍一下你自己", "用 Python 写一个快速排序", "1+1等于几？"]
]

# 4. 运行（阻塞，直到所有任务完成）
scheduler.run()

# 5. 获取结果
for task in tasks:
    print(f"[{task.task_id}] {task.decode(task.pop_tokens())[:200]}...")

loader = RWKV070ModelLoader("models/rwkv7-2.9b.pth", "eof_v20230424.txt")

# 常用方法
loader.encode("你好")                # str → list[int]
loader.decode([123, 456])            # list[int] → str
loader.gen_state(batch_size=1)       # 生成初始 state（用于 fork 多个 task）

def on_tokens(tokens):
    """每次 pulse 生成 buffer_size 个 token 后回调"""
    text = loader.decode(tokens)
    print(text, end="", flush=True)

def on_finish(all_tokens):
    """任务完成回调"""
    print(f"\n--- 完成，共 {len(all_tokens)} tokens ---")

task = scheduler.new_task(
    "User: 写一首关于春天的诗\n\nAssistant:",
    max_tokens=300,
    temperature=0.8,
    collect_callback=on_tokens,   # 每个 pulse 回调
    finish_callback=on_finish,    # 完成时回调
)
scheduler.run()

from server import RWKV070ModelLoader, DynamicScheduler
from server.task.manager import get_task_manager
from server.config import settings

loader = RWKV070ModelLoader(settings.model_path, settings.vocab_path)
scheduler = DynamicScheduler(loader, 256, 32)
task_manager = get_task_manager()
task_manager.set_dependencies(loader, scheduler.sampler)

# Step 1: 创建模板 —— max_tokens=0，只 prefill 不生成
tpl = scheduler.new_task(
    "User: 你是一个资深 Python 后端工程师，请遵循 PEP8 规范。\n\nAssistant: 好的。\n\n",
    max_tokens=0,
    temperature=0.3,
)
scheduler.run()  # prefill 完成，state 在 CPU

# Step 2: 存入 TaskManager（SQLite 持久化）
task_manager.put_task("_coder", tpl)

# Step 3: 使用时 fork —— state 自动继承，零 prefill 开销
for question in ["写一个线程安全的单例", "如何优化数据库连接池？", "解释 GIL 及其影响"]:
    new_id = task_manager.fork_template("_coder")
    task = task_manager.get_task(new_id)
    task.continue_gen()  # 重置状态为 READY
    task.prefill(f"User: {question}\n\nAssistant:")  # 追加新 prompt
    task.max_tokens = 500
    scheduler.add_task(task)

scheduler.run()  # 3 个任务并行推理，共享同一份系统提示词 prefill

# 启动后台线程
scheduler.start_daemon()  # 内部循环：有任务就 run()，空闲 sleep 0.5s

# 随时提交新任务
task = scheduler.new_task("User: 你好\n\nAssistant:", max_tokens=100)
# 调度器自动发现 READY 任务并开始推理

# 优雅关闭
scheduler.shutdown()

# FIM 本质是 prompt 构造，两行即可
prompt = f"✿prefix✿✿suffix✿{suffix}✿middle✿{prefix}"
task = scheduler.new_task(prompt, max_tokens=100, temperature=0.8)
scheduler.run()
print(loader.decode(task.pop_tokens()))

from server import RWKV070ModelLoader, DynamicScheduler
from server.config import settings

loader = RWKV070ModelLoader(settings.model_path, settings.vocab_path)
scheduler = DynamicScheduler(loader, max_batch_size=64, buffer_size=32)

def translate(texts: list[str], source_lang="中文", target_lang="英文"):
    tasks = [
        scheduler.new_task(
            f"User: 将以下{source_lang}翻译为{target_lang}：{t}\n\nAssistant:",
            max_tokens=300,
            temperature=0.3,
            top_p=0.3,
        )
        for t in texts
    ]
    scheduler.run()
    return [loader.decode(t.pop_tokens()) for t in tasks]

# 使用
results = translate(["你好世界", "机器学习很有趣", "今天天气真好"])
for r in results:
    print(r)

from rwkv_api import Client

client = Client("http://localhost:8000")

# 创建任务并等待结果 —— 服务端自动批量调度
task = client.create("Hello world", max_tokens=50)
result = task.wait()
print(result.result)

# 实时流式输出
for chunk in client.create_stream("Tell me a story", max_tokens=200):
    print(chunk, end="", flush=True)

from rwkv_api import Client

client = Client("http://localhost:8000")

# 先创建任务（非流式）
task = client.create("Hello", max_tokens=200)

# 订阅实时流式输出，重连时传 pos 避免重复
for chunk in task.stream():
    print(chunk, end="", flush=True)