playlist-admin 2.15.24 → 2.15.25
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- package/cli/index.js +22 -0
- package/package.json +2 -2
- package/rag/study_build.py +311 -0
- package/rag/study_query.py +180 -0
package/cli/index.js
CHANGED
|
@@ -21,6 +21,8 @@ Usage:
|
|
|
21
21
|
playlist-admin favorite toggle <song> Toggle favorite (我的最愛)
|
|
22
22
|
playlist-admin rag build [--reset] Build podcast RAG vector DB
|
|
23
23
|
playlist-admin rag query "問題" [--topk N] [--show 節目] [--json]
|
|
24
|
+
playlist-admin study build [--reset] Build study RAG (PDF + 課程)
|
|
25
|
+
playlist-admin study query "問題" [--topk N] [--category X] [--json]
|
|
24
26
|
`;
|
|
25
27
|
|
|
26
28
|
function projectRoot() {
|
|
@@ -121,6 +123,22 @@ function runRag(args) {
|
|
|
121
123
|
forwardPy([script, ...args.slice(1)]);
|
|
122
124
|
}
|
|
123
125
|
|
|
126
|
+
function runStudy(args) {
|
|
127
|
+
const sub = args[0];
|
|
128
|
+
if (sub !== 'build' && sub !== 'query') {
|
|
129
|
+
console.error(`未知 study 子命令: ${sub ?? '(空)'}\n用法: playlist-admin study build | playlist-admin study query "問題"`);
|
|
130
|
+
process.exit(1);
|
|
131
|
+
}
|
|
132
|
+
const roots = [projectRoot(), packageRoot()].filter(Boolean);
|
|
133
|
+
const fname = sub === 'build' ? 'study_build.py' : 'study_query.py';
|
|
134
|
+
const script = roots.map((r) => path.join(r, 'rag', fname)).find((s) => fs.existsSync(s));
|
|
135
|
+
if (!script) {
|
|
136
|
+
console.error(`找不到 rag/${fname}。請在專案內執行,或設定 PA_ROOT`);
|
|
137
|
+
process.exit(1);
|
|
138
|
+
}
|
|
139
|
+
forwardPy([script, ...args.slice(1)]);
|
|
140
|
+
}
|
|
141
|
+
|
|
124
142
|
async function main() {
|
|
125
143
|
const args = process.argv.slice(2);
|
|
126
144
|
if (args.length === 0) {
|
|
@@ -131,6 +149,10 @@ async function main() {
|
|
|
131
149
|
runRag(args.slice(1));
|
|
132
150
|
return;
|
|
133
151
|
}
|
|
152
|
+
if (args[0] === 'study') {
|
|
153
|
+
runStudy(args.slice(1));
|
|
154
|
+
return;
|
|
155
|
+
}
|
|
134
156
|
|
|
135
157
|
const root = projectRoot();
|
|
136
158
|
if (!root) {
|
package/package.json
CHANGED
|
@@ -1,7 +1,7 @@
|
|
|
1
1
|
{
|
|
2
2
|
"name": "playlist-admin",
|
|
3
|
-
"version": "2.15.
|
|
4
|
-
"description": "playlist-admin CLI -
|
|
3
|
+
"version": "2.15.25",
|
|
4
|
+
"description": "playlist-admin CLI - pipeline, podcast, music + podcast RAG + study RAG (PDF 試題/課本)",
|
|
5
5
|
"type": "module",
|
|
6
6
|
"bin": {
|
|
7
7
|
"playlist-admin": "./cli/index.js"
|
|
@@ -0,0 +1,311 @@
|
|
|
1
|
+
"""Build study RAG: PDF 試題/課本 + 課程逐字稿 → ChromaDB study collection
|
|
2
|
+
用法:
|
|
3
|
+
python rag/study_build.py # 增量建庫
|
|
4
|
+
python rag/study_build.py --reset # 清空重建
|
|
5
|
+
python rag/study_build.py --pdf-root <DIR> # 自訂 PDF 根目錄
|
|
6
|
+
python rag/study_build.py --limit 10 # 只建前 10 個 PDF(測試用)
|
|
7
|
+
"""
|
|
8
|
+
import argparse
|
|
9
|
+
import hashlib
|
|
10
|
+
import os
|
|
11
|
+
import re
|
|
12
|
+
import sys
|
|
13
|
+
import time
|
|
14
|
+
from concurrent.futures import ThreadPoolExecutor, as_completed
|
|
15
|
+
from pathlib import Path
|
|
16
|
+
|
|
17
|
+
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
|
|
18
|
+
|
|
19
|
+
import chromadb
|
|
20
|
+
import requests
|
|
21
|
+
|
|
22
|
+
from _resolve import find_base, chroma_db_dir
|
|
23
|
+
|
|
24
|
+
if hasattr(sys.stdout, "reconfigure"):
|
|
25
|
+
sys.stdout.reconfigure(encoding="utf-8")
|
|
26
|
+
|
|
27
|
+
COLLECTION = "study"
|
|
28
|
+
OLLAMA_URL = os.environ.get("OLLAMA_URL", "http://localhost:11434")
|
|
29
|
+
EMBED_MODEL = os.environ.get("OLLAMA_EMBED_MODEL", "bge-m3")
|
|
30
|
+
|
|
31
|
+
# ── Embedding ────────────────────────────────────────────────────────
|
|
32
|
+
|
|
33
|
+
def embed_one_batch(model: str, texts: list[str]) -> tuple[list[str], list[list[float]]]:
|
|
34
|
+
resp = requests.post(
|
|
35
|
+
f"{OLLAMA_URL}/api/embed",
|
|
36
|
+
json={"model": model, "input": texts},
|
|
37
|
+
timeout=120,
|
|
38
|
+
)
|
|
39
|
+
if resp.status_code == 400:
|
|
40
|
+
ok, vecs = [], []
|
|
41
|
+
for t in texts:
|
|
42
|
+
try:
|
|
43
|
+
r2 = requests.post(
|
|
44
|
+
f"{OLLAMA_URL}/api/embed",
|
|
45
|
+
json={"model": model, "input": [t]}, timeout=120,
|
|
46
|
+
)
|
|
47
|
+
r2.raise_for_status()
|
|
48
|
+
ok.append(t)
|
|
49
|
+
vecs.append(r2.json()["embeddings"][0])
|
|
50
|
+
except Exception:
|
|
51
|
+
pass
|
|
52
|
+
return ok, vecs
|
|
53
|
+
resp.raise_for_status()
|
|
54
|
+
return texts, resp.json()["embeddings"]
|
|
55
|
+
|
|
56
|
+
|
|
57
|
+
def embed_texts(model: str, texts: list[str], batch: int = 32, workers: int = 6):
|
|
58
|
+
batches = [texts[i:i + batch] for i in range(0, len(texts), batch)]
|
|
59
|
+
if len(batches) <= 1:
|
|
60
|
+
return embed_one_batch(model, batches[0]) if batches else ([], [])
|
|
61
|
+
kept, vecs = [], []
|
|
62
|
+
results = [None] * len(batches)
|
|
63
|
+
with ThreadPoolExecutor(max_workers=workers) as ex:
|
|
64
|
+
fut = {ex.submit(embed_one_batch, model, b): i for i, b in enumerate(batches)}
|
|
65
|
+
for f in as_completed(fut):
|
|
66
|
+
results[fut[f]] = f.result()
|
|
67
|
+
for k, v in results:
|
|
68
|
+
if k is None:
|
|
69
|
+
continue
|
|
70
|
+
kept.extend(k)
|
|
71
|
+
vecs.extend(v)
|
|
72
|
+
return kept, vecs
|
|
73
|
+
|
|
74
|
+
|
|
75
|
+
# ── PDF extraction ───────────────────────────────────────────────────
|
|
76
|
+
|
|
77
|
+
def extract_pdf_text(pdf_path: Path) -> str:
|
|
78
|
+
"""用 PyMuPDF 提取 PDF 文字"""
|
|
79
|
+
try:
|
|
80
|
+
import fitz
|
|
81
|
+
doc = fitz.open(str(pdf_path))
|
|
82
|
+
parts = []
|
|
83
|
+
for page in doc:
|
|
84
|
+
parts.append(page.get_text())
|
|
85
|
+
doc.close()
|
|
86
|
+
return "\n".join(parts)
|
|
87
|
+
except Exception as e:
|
|
88
|
+
return f"[PDF extraction error: {e}]"
|
|
89
|
+
|
|
90
|
+
|
|
91
|
+
# ── Chunking ─────────────────────────────────────────────────────────
|
|
92
|
+
|
|
93
|
+
def split_sentences(text: str) -> list[str]:
|
|
94
|
+
text = re.sub(r"\s+", " ", text).strip()
|
|
95
|
+
parts = re.split(r"(?<=[。!?!?;;\n])", text)
|
|
96
|
+
return [p.strip() for p in parts if len(p.strip()) >= 4]
|
|
97
|
+
|
|
98
|
+
|
|
99
|
+
def make_chunks(sentences: list[str], max_len: int = 800, overlap: int = 80) -> list[str]:
|
|
100
|
+
"""句子接成 chunk,max_len 較大(試題需要更多上下文)"""
|
|
101
|
+
chunks: list[str] = []
|
|
102
|
+
cur = ""
|
|
103
|
+
for s in sentences:
|
|
104
|
+
if len(cur) + len(s) > max_len and cur:
|
|
105
|
+
chunks.append(cur.strip())
|
|
106
|
+
# overlap: 保留尾巴
|
|
107
|
+
overlap_text = cur[-overlap:] if len(cur) > overlap else cur
|
|
108
|
+
cur = overlap_text + " " + s
|
|
109
|
+
else:
|
|
110
|
+
cur = cur + " " + s if cur else s
|
|
111
|
+
if cur.strip():
|
|
112
|
+
chunks.append(cur.strip())
|
|
113
|
+
# 超長強制硬切
|
|
114
|
+
final = []
|
|
115
|
+
for ch in chunks:
|
|
116
|
+
if len(ch) > 2800:
|
|
117
|
+
final.extend(ch[i:i + 2800] for i in range(0, len(ch), 2800))
|
|
118
|
+
else:
|
|
119
|
+
final.append(ch)
|
|
120
|
+
return final
|
|
121
|
+
|
|
122
|
+
|
|
123
|
+
def file_sig(path: Path) -> str:
|
|
124
|
+
return hashlib.md5(str(path.resolve()).encode()).hexdigest()
|
|
125
|
+
|
|
126
|
+
|
|
127
|
+
def content_sig(text: str) -> str:
|
|
128
|
+
return hashlib.md5(text[:4096].encode()).hexdigest()
|
|
129
|
+
|
|
130
|
+
|
|
131
|
+
# ── PDF discovery ────────────────────────────────────────────────────
|
|
132
|
+
|
|
133
|
+
def discover_pdfs(root: Path) -> list[tuple[str, Path]]:
|
|
134
|
+
"""回傳 (category, pdf_path) 列表"""
|
|
135
|
+
results = []
|
|
136
|
+
for dirpath, dirnames, filenames in os.walk(str(root)):
|
|
137
|
+
dirnames[:] = [d for d in dirnames if not d.startswith('.')]
|
|
138
|
+
dp = Path(dirpath)
|
|
139
|
+
rel = dp.relative_to(root)
|
|
140
|
+
category = str(rel).replace("\\", "/") if str(rel) != "." else root.name
|
|
141
|
+
for f in sorted(filenames):
|
|
142
|
+
if f.lower().endswith(".pdf"):
|
|
143
|
+
results.append((category, dp / f))
|
|
144
|
+
return results
|
|
145
|
+
|
|
146
|
+
|
|
147
|
+
def discover_txt_transcripts(podcasts_dir: Path) -> list[tuple[str, Path]]:
|
|
148
|
+
"""掃描 podcast 資料夾中的課程目錄(非 RSS)"""
|
|
149
|
+
results = []
|
|
150
|
+
# 課程目錄名單(有 .txt 逐字稿的非 RSS 資料夾)
|
|
151
|
+
course_dirs = [
|
|
152
|
+
'T夜工智CoCo英文S', 'T夜工智張瑀國文S', 'T夜工智胡傑數學S',
|
|
153
|
+
'機械群劉徹機械力學', '機械群劉徹機械力學S',
|
|
154
|
+
'機械群許凱機件原理', '機械群許凱機件原理S',
|
|
155
|
+
'機械群陳海程機械製圖', '機械群陳海程機械製圖S',
|
|
156
|
+
'機械群陳海程機械製造與實習', '機械群陳海程機械製造與實習S',
|
|
157
|
+
]
|
|
158
|
+
for name in course_dirs:
|
|
159
|
+
d = podcasts_dir / name
|
|
160
|
+
if not d.is_dir():
|
|
161
|
+
continue
|
|
162
|
+
for f in sorted(d.glob("*.txt")):
|
|
163
|
+
results.append((f"課程/{name}", f))
|
|
164
|
+
return results
|
|
165
|
+
|
|
166
|
+
|
|
167
|
+
# ── Main ─────────────────────────────────────────────────────────────
|
|
168
|
+
|
|
169
|
+
def main():
|
|
170
|
+
ap = argparse.ArgumentParser(description="Build study RAG")
|
|
171
|
+
ap.add_argument("--reset", action="store_true", help="清空重建")
|
|
172
|
+
ap.add_argument("--pdf-root", default="", help="PDF 根目錄")
|
|
173
|
+
ap.add_argument("--limit", type=int, default=0, help="只建前 N 個 PDF(測試)")
|
|
174
|
+
ap.add_argument("--db", default=str(chroma_db_dir()))
|
|
175
|
+
ap.add_argument("--embed", default=EMBED_MODEL)
|
|
176
|
+
ap.add_argument("--workers", type=int, default=6)
|
|
177
|
+
args = ap.parse_args()
|
|
178
|
+
|
|
179
|
+
# Resolve PDF root
|
|
180
|
+
if args.pdf_root:
|
|
181
|
+
pdf_root = Path(args.pdf_root)
|
|
182
|
+
else:
|
|
183
|
+
# Auto-discover: check Desktop, then common locations
|
|
184
|
+
home = Path.home()
|
|
185
|
+
candidates = [
|
|
186
|
+
home / "Desktop" / "四技二專歷屆試題與解答",
|
|
187
|
+
find_base() / "study_materials",
|
|
188
|
+
]
|
|
189
|
+
pdf_root = None
|
|
190
|
+
for c in candidates:
|
|
191
|
+
if c.exists():
|
|
192
|
+
pdf_root = c
|
|
193
|
+
break
|
|
194
|
+
if pdf_root is None:
|
|
195
|
+
print("找不到 PDF 目錄。請用 --pdf-root 指定。")
|
|
196
|
+
sys.exit(1)
|
|
197
|
+
|
|
198
|
+
print(f"PDF 根目錄: {pdf_root}")
|
|
199
|
+
|
|
200
|
+
# Connect ChromaDB
|
|
201
|
+
client = chromadb.PersistentClient(path=args.db)
|
|
202
|
+
if args.reset:
|
|
203
|
+
try:
|
|
204
|
+
client.delete_collection(COLLECTION)
|
|
205
|
+
print("已清空 study collection")
|
|
206
|
+
except Exception:
|
|
207
|
+
pass
|
|
208
|
+
col = client.get_or_create_collection(
|
|
209
|
+
COLLECTION,
|
|
210
|
+
metadata={"hnsw:space": "cosine"},
|
|
211
|
+
)
|
|
212
|
+
|
|
213
|
+
# Existing signatures (dedup)
|
|
214
|
+
existing = col.get(include=["metadatas"])
|
|
215
|
+
done_sigs: set[str] = set()
|
|
216
|
+
if existing["ids"]:
|
|
217
|
+
for meta in existing["metadatas"]:
|
|
218
|
+
sig = meta.get("sig", "")
|
|
219
|
+
if sig:
|
|
220
|
+
done_sigs.add(sig)
|
|
221
|
+
print(f"已存在 {len(existing['ids'])} 筆, 已知 sig {len(done_sigs)} 個")
|
|
222
|
+
|
|
223
|
+
# Discover files
|
|
224
|
+
pdfs = discover_pdfs(pdf_root)
|
|
225
|
+
transcripts = discover_txt_transcripts(find_base() / "podcasts")
|
|
226
|
+
all_files = [(cat, p, "pdf") for cat, p in pdfs] + [(cat, p, "txt") for cat, p in transcripts]
|
|
227
|
+
print(f"發現 {len(pdfs)} 個 PDF, {len(transcripts)} 個逐字稿")
|
|
228
|
+
|
|
229
|
+
if args.limit:
|
|
230
|
+
all_files = all_files[:args.limit]
|
|
231
|
+
print(f"測試模式: 只建前 {args.limit} 個")
|
|
232
|
+
|
|
233
|
+
# Extract → chunk → embed
|
|
234
|
+
total_chunks = 0
|
|
235
|
+
t_start = time.time()
|
|
236
|
+
|
|
237
|
+
for fi, (category, fpath, ftype) in enumerate(all_files, 1):
|
|
238
|
+
if ftype == "pdf":
|
|
239
|
+
sig = file_sig(fpath)
|
|
240
|
+
if sig in done_sigs:
|
|
241
|
+
continue
|
|
242
|
+
text = extract_pdf_text(fpath)
|
|
243
|
+
else:
|
|
244
|
+
sig = file_sig(fpath)
|
|
245
|
+
if sig in done_sigs:
|
|
246
|
+
continue
|
|
247
|
+
text = fpath.read_text(encoding="utf-8-sig", errors="replace")
|
|
248
|
+
|
|
249
|
+
if not text or len(text.strip()) < 20:
|
|
250
|
+
continue
|
|
251
|
+
|
|
252
|
+
chunks = make_chunks(split_sentences(text))
|
|
253
|
+
if not chunks:
|
|
254
|
+
continue
|
|
255
|
+
|
|
256
|
+
# Token safety: 2800 chars per chunk
|
|
257
|
+
final_chunks = []
|
|
258
|
+
for ch in chunks:
|
|
259
|
+
if len(ch) > 2800:
|
|
260
|
+
final_chunks.extend(ch[i:i + 2800] for i in range(0, len(ch), 2800))
|
|
261
|
+
else:
|
|
262
|
+
final_chunks.append(ch)
|
|
263
|
+
chunks = final_chunks
|
|
264
|
+
|
|
265
|
+
# Embed
|
|
266
|
+
kept, vectors = embed_texts(args.embed, chunks, workers=args.workers)
|
|
267
|
+
if not kept:
|
|
268
|
+
continue
|
|
269
|
+
|
|
270
|
+
# Write to ChromaDB
|
|
271
|
+
prefix = f"study_{sig[:8]}"
|
|
272
|
+
ids = [f"{prefix}_{i}" for i in range(len(kept))]
|
|
273
|
+
metas = [
|
|
274
|
+
{
|
|
275
|
+
"category": category,
|
|
276
|
+
"file": fpath.name,
|
|
277
|
+
"path": str(fpath),
|
|
278
|
+
"sig": sig if i == 0 else "", # only first chunk carries sig
|
|
279
|
+
"chunk_index": i,
|
|
280
|
+
"total_chunks": len(kept),
|
|
281
|
+
"source_type": ftype,
|
|
282
|
+
}
|
|
283
|
+
for i in range(len(kept))
|
|
284
|
+
]
|
|
285
|
+
|
|
286
|
+
batch_size = 4500
|
|
287
|
+
for bi in range(0, len(ids), batch_size):
|
|
288
|
+
col.add(
|
|
289
|
+
ids=ids[bi:bi + batch_size],
|
|
290
|
+
documents=kept[bi:bi + batch_size],
|
|
291
|
+
metadatas=metas[bi:bi + batch_size],
|
|
292
|
+
embeddings=vectors[bi:bi + batch_size],
|
|
293
|
+
)
|
|
294
|
+
|
|
295
|
+
total_chunks += len(kept)
|
|
296
|
+
elapsed = time.time() - t_start
|
|
297
|
+
pct = fi / len(all_files) * 100
|
|
298
|
+
rate = fi / elapsed if elapsed > 0 else 0
|
|
299
|
+
eta = (len(all_files) - fi) / rate if rate > 0 else 0
|
|
300
|
+
print(
|
|
301
|
+
f"\r [{fi}/{len(all_files)}] {pct:.1f}% "
|
|
302
|
+
f"| {rate:.2f} files/s | ETA {eta / 60:.1f}min "
|
|
303
|
+
f"| chunks累計 {total_chunks} | {category}/{fpath.name[:30]}",
|
|
304
|
+
end="", flush=True,
|
|
305
|
+
)
|
|
306
|
+
|
|
307
|
+
print(f"\n完成! 共 {total_chunks} 個新 chunk, 總計 {col.count()} 筆")
|
|
308
|
+
|
|
309
|
+
|
|
310
|
+
if __name__ == "__main__":
|
|
311
|
+
main()
|
|
@@ -0,0 +1,180 @@
|
|
|
1
|
+
"""Study RAG 查詢: 問題 → bge-m3 → top-k 片段
|
|
2
|
+
用法:
|
|
3
|
+
python rag/study_query.py "問題" # 文字輸出
|
|
4
|
+
python rag/study_query.py "問題" --json # JSON 輸出
|
|
5
|
+
python rag/study_query.py "問題" --topk 5 # 前 5 個結果
|
|
6
|
+
python rag/study_query.py "問題" --category 機械 # 只搜某分類
|
|
7
|
+
python rag/study_query.py "問題" --answer # Ollama 生成回答
|
|
8
|
+
"""
|
|
9
|
+
import argparse
|
|
10
|
+
import json
|
|
11
|
+
import os
|
|
12
|
+
import re
|
|
13
|
+
import sys
|
|
14
|
+
from pathlib import Path
|
|
15
|
+
|
|
16
|
+
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
|
|
17
|
+
|
|
18
|
+
import chromadb
|
|
19
|
+
import requests
|
|
20
|
+
|
|
21
|
+
from _resolve import chroma_db_dir
|
|
22
|
+
|
|
23
|
+
if hasattr(sys.stdout, "reconfigure"):
|
|
24
|
+
sys.stdout.reconfigure(encoding="utf-8")
|
|
25
|
+
|
|
26
|
+
OLLAMA_URL = os.environ.get("OLLAMA_URL", "http://localhost:11434")
|
|
27
|
+
CHAT_MODEL = os.environ.get("OLLAMA_CHAT_MODEL", "qwen2.5:7b")
|
|
28
|
+
COLLECTION = "study"
|
|
29
|
+
|
|
30
|
+
|
|
31
|
+
def embed_one(model: str, text: str) -> list[float]:
|
|
32
|
+
resp = requests.post(
|
|
33
|
+
f"{OLLAMA_URL}/api/embed", json={"model": model, "input": [text]}, timeout=120
|
|
34
|
+
)
|
|
35
|
+
resp.raise_for_status()
|
|
36
|
+
return resp.json()["embeddings"][0]
|
|
37
|
+
|
|
38
|
+
|
|
39
|
+
def read_file(path_str: str) -> str:
|
|
40
|
+
p = Path(path_str)
|
|
41
|
+
if not p.exists():
|
|
42
|
+
return ""
|
|
43
|
+
for enc in ("utf-8-sig", "utf-8", "big5", "gb18030"):
|
|
44
|
+
try:
|
|
45
|
+
return p.read_text(encoding=enc)
|
|
46
|
+
except (UnicodeDecodeError, UnicodeError):
|
|
47
|
+
continue
|
|
48
|
+
return p.read_text(encoding="utf-8", errors="replace")
|
|
49
|
+
|
|
50
|
+
|
|
51
|
+
def pick_chat_models() -> list[str]:
|
|
52
|
+
forced = os.environ.get("OLLAMA_CHAT_MODEL", "").strip()
|
|
53
|
+
try:
|
|
54
|
+
tags = requests.get(f"{OLLAMA_URL}/api/tags", timeout=30).json().get("models", [])
|
|
55
|
+
except Exception:
|
|
56
|
+
return [forced] if forced else []
|
|
57
|
+
local = [m["name"] for m in tags if m.get("size", 0) > 0]
|
|
58
|
+
cloud = [m["name"] for m in tags if m.get("size", 0) == 0]
|
|
59
|
+
banned = {"bge-m3", "bge-m3:latest"}
|
|
60
|
+
local = [m for m in local if m not in banned]
|
|
61
|
+
if forced:
|
|
62
|
+
return [forced, *local, *cloud]
|
|
63
|
+
return [*local, *cloud]
|
|
64
|
+
|
|
65
|
+
|
|
66
|
+
def generate_answer(question: str, hits: list[dict], quiet: bool = False) -> str:
|
|
67
|
+
context = "\n".join(
|
|
68
|
+
f"[{h['category']} | {h['file']}]\n{h['chunk']}" for h in hits
|
|
69
|
+
)
|
|
70
|
+
prompt = (
|
|
71
|
+
"你是學習助理,專精四技二專統一入學測驗與機械群專業科目。"
|
|
72
|
+
"請只根據以下學習資料片段回答使用者的問題,"
|
|
73
|
+
"回答用繁體中文,簡潔、分點陳述,並在結尾列出引用來源。\n\n"
|
|
74
|
+
f"片段資料:\n{context}\n\n問題:{question}"
|
|
75
|
+
)
|
|
76
|
+
errors = []
|
|
77
|
+
for model in pick_chat_models():
|
|
78
|
+
try:
|
|
79
|
+
resp = requests.post(
|
|
80
|
+
f"{OLLAMA_URL}/api/chat",
|
|
81
|
+
json={"model": model, "messages": [{"role": "user", "content": prompt}], "stream": True},
|
|
82
|
+
timeout=600, stream=True,
|
|
83
|
+
)
|
|
84
|
+
if resp.status_code in (400, 401, 404):
|
|
85
|
+
errors.append(f"{model}: HTTP {resp.status_code}")
|
|
86
|
+
continue
|
|
87
|
+
resp.raise_for_status()
|
|
88
|
+
answer = []
|
|
89
|
+
for line in resp.iter_lines():
|
|
90
|
+
if not line:
|
|
91
|
+
continue
|
|
92
|
+
try:
|
|
93
|
+
chunk = json.loads(line)
|
|
94
|
+
token = chunk.get("message", {}).get("content", "")
|
|
95
|
+
if token:
|
|
96
|
+
answer.append(token)
|
|
97
|
+
if not quiet:
|
|
98
|
+
print(token, end="", flush=True)
|
|
99
|
+
except json.JSONDecodeError:
|
|
100
|
+
pass
|
|
101
|
+
if not quiet:
|
|
102
|
+
print()
|
|
103
|
+
return "".join(answer).strip()
|
|
104
|
+
except requests.RequestException as e:
|
|
105
|
+
errors.append(f"{model}: {e}")
|
|
106
|
+
if errors:
|
|
107
|
+
raise RuntimeError(" | ".join(errors))
|
|
108
|
+
return ""
|
|
109
|
+
|
|
110
|
+
|
|
111
|
+
def main():
|
|
112
|
+
ap = argparse.ArgumentParser(description="Study RAG query")
|
|
113
|
+
ap.add_argument("question", nargs="?", help="要查的問題")
|
|
114
|
+
ap.add_argument("--db", default=str(chroma_db_dir()))
|
|
115
|
+
ap.add_argument("--embed", default="bge-m3")
|
|
116
|
+
ap.add_argument("--topk", type=int, default=5)
|
|
117
|
+
ap.add_argument("--json", action="store_true")
|
|
118
|
+
ap.add_argument("--category", default="", help="只搜某分類")
|
|
119
|
+
ap.add_argument("--answer", action="store_true", help="用 Ollama 生成回答")
|
|
120
|
+
ap.add_argument("--topk-answer", type=int, default=4)
|
|
121
|
+
args = ap.parse_args()
|
|
122
|
+
|
|
123
|
+
if not args.question:
|
|
124
|
+
args.question = input("問題: ").strip()
|
|
125
|
+
if not args.question:
|
|
126
|
+
sys.exit(1)
|
|
127
|
+
|
|
128
|
+
client = chromadb.PersistentClient(path=args.db)
|
|
129
|
+
try:
|
|
130
|
+
col = client.get_collection(COLLECTION)
|
|
131
|
+
except Exception:
|
|
132
|
+
print(f"找不到 {COLLECTION} collection。請先執行: python rag/study_build.py")
|
|
133
|
+
sys.exit(1)
|
|
134
|
+
|
|
135
|
+
qvec = embed_one(args.embed, args.question)
|
|
136
|
+
fetch_n = max(args.topk * 5, 30) if args.category else args.topk
|
|
137
|
+
where_filter = {"category": {"$contains": args.category}} if args.category else None
|
|
138
|
+
res = col.query(
|
|
139
|
+
query_embeddings=[qvec],
|
|
140
|
+
n_results=fetch_n,
|
|
141
|
+
where=where_filter,
|
|
142
|
+
)
|
|
143
|
+
|
|
144
|
+
raw_hits = []
|
|
145
|
+
docs = res["documents"][0] if res.get("documents") else []
|
|
146
|
+
metas = res["metadatas"][0] if res.get("metadatas") else []
|
|
147
|
+
dists = res["distances"][0] if res.get("distances") else []
|
|
148
|
+
for doc, meta, dist in zip(docs, metas, dists):
|
|
149
|
+
if args.category and args.category not in meta.get("category", ""):
|
|
150
|
+
continue
|
|
151
|
+
raw_hits.append({
|
|
152
|
+
"similarity": round(1 - dist, 3),
|
|
153
|
+
"category": meta.get("category", "?"),
|
|
154
|
+
"file": meta.get("file", "?"),
|
|
155
|
+
"path": meta.get("path", ""),
|
|
156
|
+
"source_type": meta.get("source_type", "?"),
|
|
157
|
+
"chunk": doc,
|
|
158
|
+
})
|
|
159
|
+
|
|
160
|
+
if args.json:
|
|
161
|
+
payload = {"question": args.question, "results": raw_hits[:args.topk]}
|
|
162
|
+
if args.answer:
|
|
163
|
+
try:
|
|
164
|
+
payload["answer"] = generate_answer(
|
|
165
|
+
args.question, raw_hits[:args.topk_answer], quiet=True
|
|
166
|
+
)
|
|
167
|
+
except Exception as e:
|
|
168
|
+
payload["answer_error"] = str(e)
|
|
169
|
+
print(json.dumps(payload, ensure_ascii=False, indent=2))
|
|
170
|
+
return
|
|
171
|
+
|
|
172
|
+
print(f"\n🔍 {args.question}\n")
|
|
173
|
+
for i, h in enumerate(raw_hits[:args.topk], 1):
|
|
174
|
+
print(f"[{i}] 類似度 {h['similarity']} | {h['source_type']} | {h['category']}/{h['file']}")
|
|
175
|
+
print(f" {h['chunk'][:150]}...")
|
|
176
|
+
print()
|
|
177
|
+
|
|
178
|
+
|
|
179
|
+
if __name__ == "__main__":
|
|
180
|
+
main()
|