"""RepoBot v2：会查 httpx 官方文档的答疑助手。
和 v1 相比，每次回答之前先做三件事：把问题改写成英文检索词，在文档里检索最相关的 5 块，
把它们编号放进提示词。回答里用 [编号] 注明出处，最后列出引用到的文档。
python repobot.py 开始对话，输入空行退出
python repobot.py --rerank 检索时加上重排（更准，每个问题多约 1 秒）
python repobot.py --show-query 显示改写后的检索词和检索到的文档
python repobot.py --think 开启思考模式
"""
import os
import re
import sys
from pathlib import Path
import openai
import llm
from retrieval import QueryRewriter, Retriever
HERE = Path(__file__).parent
DOCS_DIR = os.environ.get("REPOBOT_DOCS", HERE / "../../../data/httpx-docs")
CACHE_DIR = HERE / ".cache"
MAX_MESSAGES = 20
SYSTEM = """你是 RepoBot，Python HTTP 客户端库 httpx 的答疑助手。
每个问题都会附上从 httpx 官方文档里检索到的片段，放在 <docs> 里，每个片段有编号。
- 只根据这些片段回答。每句话后面用 [编号] 注明依据，如 [2] 或 [1][3]。
- 片段里没有的信息不要写，哪怕你自己知道。片段不足以回答时，直接说"文档里没有找到相关说明"，
可以建议用户去哪里查。
- 和 httpx 无关的问题，礼貌地说明你只负责 httpx，不要回答。
- 用中文回答，简洁，代码保持原样。"""
def build_context(results):
parts = [f'<doc id="{i}" source="{file}">\n{text}\n</doc>' for i, (file, text) in enumerate(results, 1)]
return "<docs>\n" + "\n".join(parts) + "\n</docs>"
def stream_answer(messages, thinking):
parts, usage, finish = [], None, None
for chunk in llm.open_stream(messages, thinking):
if chunk.usage:
usage = chunk.usage
if not chunk.choices:
continue
choice = chunk.choices[0]
if choice.delta.content:
parts.append(choice.delta.content)
print(choice.delta.content, end="", flush=True)
if choice.finish_reason:
finish = choice.finish_reason
print()
if finish == "length":
print("[回答太长，被截断了]")
return "".join(parts), usage
def main():
thinking, show = "--think" in sys.argv, "--show-query" in sys.argv
print("正在加载文档和模型……", flush=True)
retriever = Retriever(DOCS_DIR, CACHE_DIR, rerank="--rerank" in sys.argv)
rewriter = QueryRewriter(CACHE_DIR)
print(f"RepoBot v2（{llm.MODEL}，{len(retriever.chunks)} 个文档块）。问我 httpx 的问题，输入空行退出。")
history, total = [], 0.0
interactive = sys.stdin.isatty()
while True:
try:
question = input("\n你：" if interactive else "").strip()
except EOFError:
break
if not question:
break
if not interactive:
print(f"\n你：{question}")
try:
query = rewriter.rewrite(question, history)
results = retriever.search(query, k=5)
if show:
print(f"[检索词] {query}")
print("[检索到] " + " ".join(f"[{i}] {f}" for i, (f, _) in enumerate(results, 1)))
# 文档只放进这一轮的消息里；存进历史的只有原始问题，避免历史越来越长
messages = ([{"role": "system", "content": SYSTEM}] + history +
[{"role": "user", "content": build_context(results) + f"\n\n问题：{question}"}])
print("RepoBot：", end="", flush=True)
text, usage = stream_answer(messages, thinking)
except openai.APIError as e:
print(f"\n[出错了：{type(e).__name__}，这一轮作废，可以再问一次]")
continue
cited = sorted({int(n) for n in re.findall(r"\[(\d+)\]", text) if 1 <= int(n) <= len(results)})
if cited:
print("来源：" + " ".join(f"[{n}] {results[n - 1][0]}" for n in cited))
cost = llm.cost_usd(usage) + llm.cost_usd(rewriter.last_usage)
total += cost
print(f"[回答输入 {usage.prompt_tokens}（缓存命中 {getattr(usage, 'prompt_cache_hit_tokens', 0) or 0}），"
f"输出 {usage.completion_tokens}，本轮 {cost:.5f} 美元（含改写），累计 {total:.5f} 美元]")
history += [{"role": "user", "content": question}, {"role": "assistant", "content": text}]
history = history[-MAX_MESSAGES:]
if __name__ == "__main__":
main()