rag_core.py

from pathlib import Path
from openai import OpenAI

# ===== LLM client =====
iflow_client = OpenAI(
    base_url="https://apis.iflow.cn/v1",
    api_key="c194b99ec60f6b5682a2a9862c7356db"
)
def call_llm(prompt: str) -> str:
    resp = iflow_client.chat.completions.create(
        model="iflow-rome-30ba3b",
        messages=[
            {"role": "user", "content": prompt}
        ],
        temperature=0.2 #
    )

    if not resp.choices:
        return "模型未返回有效结果。"

    return resp.choices[0].message.content.strip()
def load_text(path: str) -> str:
    return Path(path).read_text(encoding="utf-8")

def chunk_text(text: str, chunk_size: int = 200):
    chunks = []
    for i in range(0, len(text), chunk_size):
        chunk = text[i:i+chunk_size].strip()
        if chunk:
            chunks.append({
                "chunk_id": f"c{i//chunk_size:03d}",
                "text": chunk
            })
    return chunks

#
def retrieve(query: str, chunks: list[dict], top_k: int = 3):
    q_words = list(query)   # 👈 每个字一个 token
    scored = []
    for ch in chunks:
        score = sum(ch["text"].count(w) for w in q_words)
        scored.append((score, ch))
    scored.sort(key=lambda x: x[0], reverse=True)
    return [ch for score, ch in scored[:top_k] if score > 0]

def answer_with_llm_and_citations(query: str, hits: list[dict]) -> dict:    #
    if not hits:    #
        return {"answer": "未找到依据。", "citations": []}

    context = "\n".join([f"[{h['chunk_id']}] {h['text']}" for h in hits])   #

    prompt = f"""
只能基于证据回答，不要编造。

问题：{query}

证据：
{context}

返回 JSON：
{{"answer": "...", "citations": ["c000"]}}
"""

    resp = iflow_client.chat.completions.create(     # 👈 调用 iflow
        model="iflow-rome-30ba3b",
        messages=[{"role": "user", "content": prompt}],
        temperature=0 # 👈 温度设为 0 是为了让模型更确定，避免随机输出#
    )

    import json
    try:
        return json.loads(resp.choices[0].message.content)  #
    except Exception:
        return {    # 兜底
            "answer": resp.choices[0].message.content,
            "citations": [h["chunk_id"] for h in hits]
        }