import os

from llama_index.core import Document, SummaryIndex
from llama_index.core.llms import ChatMessage
from llama_index.core.node_parser import SentenceSplitter
from llama_index.llms.openai_like import OpenAILike


def main():
    llm = OpenAILike(
        model=os.getenv("QUICKROUTER_MODEL", "glm-5.3"),
        api_base=os.getenv("QUICKROUTER_BASE_URL", "https://api.quickrouter.ai/v1"),
        api_key=os.environ["QUICKROUTER_API_KEY"],
        is_chat_model=True,
        is_function_calling_model=False,
        context_window=int(os.getenv("QUICKROUTER_CONTEXT_TOKENS", "32000")),
        max_tokens=1024,
        timeout=60,
        max_retries=0,
    )
    documents = [
        Document(text="[policy-01] 示例商店：未使用的商品可在收货后 7 天内申请退货。", metadata={"source_id": "policy-01"}),
        Document(text="[policy-02] 示例商店：定制商品不适用上述退货规则；质量问题联系人工客服。", metadata={"source_id": "policy-02"}),
    ]
    index = SummaryIndex.from_documents(
        documents,
        transformations=[SentenceSplitter(chunk_size=512, chunk_overlap=32)],
    )
    question = "定制商品能在 7 天内无理由退货吗？"
    candidates = index.as_retriever().retrieve(question)
    context = "\n".join(node.node.get_content() for node in candidates)
    result = llm.chat([
        ChatMessage(role="system", content="只根据资料回答。资料不足时说无法确认，答案引用 [source_id]。"),
        ChatMessage(role="user", content=f"资料：\n{context}\n问题：{question}"),
    ])
    answer = result.message.content
    raw = result.raw.model_dump()
    finish_reason = raw["choices"][0]["finish_reason"]
    print("answer:", answer)
    print("source_ids:", [node.node.metadata.get("source_id") for node in candidates])
    print("finish_reason:", finish_reason)
    print("usage:", raw.get("usage"))
    if not isinstance(answer, str) or not answer.strip():
        raise RuntimeError("答案为空；先检查 API 记录与模型配置。")
    if finish_reason == "length":
        raise RuntimeError("输出被截断；先核对账单，再调整输出预算。")


if __name__ == "__main__":
    main()