feat: 升级 AI 索引构建与检索系统,支持向量嵌入和混合搜索

- 重构 `build_ai_index.py`,新增向量嵌入支持,包括批量嵌入处理、API 响应校验和超时优化,提升索引生成稳定性
- 引入 `retrieve_ai_context.py` 的向量检索能力,支持基于余弦相似度的混合关键词+语义搜索,显著提高相关性匹配精度
- 更新索引清单格式至 version 2,增加生成时间、扫描文件数、启用嵌入等元数据字段,便于调试与监控
- 优化工作流 `ai-build-index.yml`,添加并发控制、步骤摘要输出和嵌入开关逻辑,增强可观测性与资源管理
- 改进 `ai-issue-smart-reply.yml` 的上下文获取流程,统一最终上下文与候选信息处理路径,确保数据一致性和容错能力
This commit is contained in:
Abner
2026-04-13 20:17:38 +08:00
parent 142c484783
commit 6c20b62670
5 changed files with 570 additions and 117 deletions

View File

@@ -7,7 +7,9 @@ import os
import pathlib
import re
import sys
import time
import urllib.request
from typing import Dict, List, Tuple
TEXT_EXTS = {
".md", ".mdx", ".rst", ".txt", ".yml", ".yaml", ".json", ".toml", ".ini", ".env",
@@ -20,21 +22,30 @@ IGNORE_DIRS = {
"__pycache__", ".venv", "venv", "vendor", "target", "out"
}
MAX_TEXT_PER_CHUNK = 3200
MAX_SUMMARY_CHARS = 900
EMBED_BATCH_SIZE = 32
def sha1(s: str) -> str:
return hashlib.sha1(s.encode("utf-8")).hexdigest()
def read_text(path: pathlib.Path) -> str:
try:
return path.read_text(encoding="utf-8", errors="ignore")
except Exception:
return ""
def tokenize(text: str):
def tokenize(text: str) -> List[str]:
return re.findall(r"[A-Za-z0-9_./:#-]{3,}", text.lower())
def file_weight(path: str) -> int:
p = path.lower()
score = 1
if "readme" in p or "changelog" in p or "/docs/" in p or p.startswith("docs/"):
score += 5
if ".env.example" in p or "config" in p or "docker-compose" in p:
@@ -45,30 +56,40 @@ def file_weight(path: str) -> int:
score += 3
if "/app/" in p or p.startswith("app/"):
score += 2
if "/lib/" in p or p.startswith("lib/"):
score += 2
if "/tests/" in p or p.startswith("tests/"):
score += 1
if "/.github/" in p or p.startswith(".github/"):
score += 2
return score
def iter_files(root: pathlib.Path, max_files: int):
count = 0
for base, dirs, files in os.walk(root):
dirs[:] = [d for d in dirs if d not in IGNORE_DIRS]
for f in files:
for f in sorted(files):
path = pathlib.Path(base) / f
rel = path.relative_to(root).as_posix()
if path.suffix.lower() not in TEXT_EXTS and not f.startswith("README"):
if path.suffix.lower() not in TEXT_EXTS and not f.startswith("README") and not f.startswith("CHANGELOG"):
continue
yield path, rel
count += 1
if count >= max_files:
return
def chunk_lines(text: str, max_chunk_chars: int):
def chunk_lines(text: str, max_chunk_chars: int) -> List[Tuple[int, int, str]]:
lines = text.splitlines()
chunks = []
buf = []
start = 1
cur = 0
for i, line in enumerate(lines, start=1):
line_len = len(line) + 1
if buf and cur + line_len > max_chunk_chars:
@@ -81,21 +102,40 @@ def chunk_lines(text: str, max_chunk_chars: int):
start = i
buf.append(line)
cur += line_len
if buf:
chunks.append((start, len(lines), "\n".join(buf)))
return chunks
def summarize_chunk(rel, chunk):
header = rel
first_lines = "\n".join(chunk.splitlines()[:12]).strip()
return f"{header}\n{first_lines}"[:800]
def embed_texts(base_url, api_key, model, texts):
def summarize_chunk(rel: str, chunk: str) -> str:
lines = chunk.splitlines()
head = "\n".join(lines[:12]).strip()
text = f"{rel}\n{head}".strip()
return text[:MAX_SUMMARY_CHARS]
def build_keywords(text: str, limit: int = 120) -> List[str]:
toks = tokenize(text)
seen = set()
out = []
for t in toks:
if t not in seen:
seen.add(t)
out.append(t)
if len(out) >= limit:
break
return out
def embed_texts(base_url: str, api_key: str, model: str, texts: List[str]) -> List[List[float]]:
url = base_url.rstrip("/") + "/embeddings"
payload = json.dumps({
"model": model,
"input": texts
}).encode("utf-8")
req = urllib.request.Request(
url,
data=payload,
@@ -105,9 +145,22 @@ def embed_texts(base_url, api_key, model, texts):
},
method="POST"
)
with urllib.request.urlopen(req, timeout=120) as resp:
with urllib.request.urlopen(req, timeout=180) as resp:
data = json.loads(resp.read().decode("utf-8"))
return [item["embedding"] for item in data["data"]]
if "data" not in data or not isinstance(data["data"], list):
raise RuntimeError("Invalid embedding API response: missing data list")
vectors = []
for item in data["data"]:
vec = item.get("embedding")
if not isinstance(vec, list):
raise RuntimeError("Invalid embedding API response: missing embedding")
vectors.append(vec)
return vectors
def main():
ap = argparse.ArgumentParser()
@@ -125,33 +178,43 @@ def main():
root = pathlib.Path(args.repo_root).resolve()
out_path = pathlib.Path(args.out)
out_path.parent.mkdir(parents=True, exist_ok=True)
manifest_path = pathlib.Path(args.manifest)
manifest_path.parent.mkdir(parents=True, exist_ok=True)
rows = []
embedding_rows = []
rows: List[Dict] = []
scanned_files = 0
kept_files = 0
for path, rel in iter_files(root, args.max_files):
scanned_files += 1
text = read_text(path)
if not text.strip():
continue
chunks = chunk_lines(text, args.max_chunk_chars)
if not chunks:
continue
kept_files += 1
weight = file_weight(rel)
for start, end, chunk in chunks:
tokens = tokenize(chunk)
if len(tokens) < 3:
kws = build_keywords(chunk)
if len(kws) < 3:
continue
text_for_store = chunk[:MAX_TEXT_PER_CHUNK]
row = {
"id": sha1(f"{rel}:{start}:{end}:{sha1(chunk)}"),
"id": sha1(f"{rel}:{start}:{end}:{sha1(text_for_store)}"),
"path": rel,
"start_line": start,
"end_line": end,
"weight": weight,
"token_count": len(tokens),
"keywords": sorted(list(set(tokens[:120]))),
"token_count": len(kws),
"keywords": kws,
"summary": summarize_chunk(rel, chunk),
"text": chunk[:3000]
"text": text_for_store
}
rows.append(row)
@@ -160,21 +223,35 @@ def main():
f.write(json.dumps(row, ensure_ascii=False) + "\n")
manifest = {
"version": 1,
"version": 2,
"generated_at_unix": int(time.time()),
"repo_root": str(root),
"scanned_files": scanned_files,
"indexed_files": kept_files,
"chunks": len(rows),
"max_files": args.max_files,
"max_chunk_chars": args.max_chunk_chars
"max_chunk_chars": args.max_chunk_chars,
"embedding_enabled": bool(args.embedding_out and args.embedding_api_key and args.embedding_base_url and args.embedding_model)
}
pathlib.Path(args.manifest).write_text(
manifest_path.write_text(
json.dumps(manifest, ensure_ascii=False, indent=2),
encoding="utf-8"
)
if args.embedding_out and args.embedding_api_key and args.embedding_base_url and args.embedding_model:
emb_out = pathlib.Path(args.embedding_out)
emb_out.parent.mkdir(parents=True, exist_ok=True)
texts = []
metas = []
for row in rows:
texts.append(f"{row['path']} lines {row['start_line']}-{row['end_line']}\n{row['summary']}\n{row['text'][:1200]}")
emb_text = (
f"{row['path']} lines {row['start_line']}-{row['end_line']}\n"
f"{row['summary']}\n"
f"{row['text'][:1200]}"
)
texts.append(emb_text)
metas.append({
"id": row["id"],
"path": row["path"],
@@ -183,16 +260,20 @@ def main():
"weight": row["weight"]
})
emb_out = pathlib.Path(args.embedding_out)
with emb_out.open("w", encoding="utf-8") as f:
batch_size = 32
for i in range(0, len(texts), batch_size):
batch = texts[i:i+batch_size]
vecs = embed_texts(args.embedding_base_url, args.embedding_api_key, args.embedding_model, batch)
for meta, vec in zip(metas[i:i+batch_size], vecs):
for i in range(0, len(texts), EMBED_BATCH_SIZE):
batch = texts[i:i + EMBED_BATCH_SIZE]
vecs = embed_texts(
args.embedding_base_url,
args.embedding_api_key,
args.embedding_model,
batch
)
for meta, vec in zip(metas[i:i + EMBED_BATCH_SIZE], vecs):
item = dict(meta)
item["embedding"] = vec
f.write(json.dumps(item, ensure_ascii=False) + "\n")
if __name__ == "__main__":
main()

View File

@@ -2,39 +2,118 @@
import argparse
import json
import math
import os
import pathlib
import re
from collections import Counter
import urllib.request
from typing import Dict, List, Tuple
def tokenize(text: str):
TEXT_EXTS = {
".md", ".mdx", ".rst", ".txt", ".js", ".jsx", ".ts", ".tsx", ".py", ".go", ".java",
".rs", ".php", ".rb", ".yml", ".yaml", ".json", ".toml", ".sh", ".vue"
}
def tokenize(text: str) -> List[str]:
return re.findall(r"[A-Za-z0-9_./:#-]{3,}", text.lower())
def load_json(path):
def load_json(path: str) -> Dict:
return json.loads(pathlib.Path(path).read_text(encoding="utf-8"))
def score_keywords(query_tokens, keywords, weight=1):
def score_keywords(query_tokens: set, keywords: List[str], weight: int = 1) -> float:
if not query_tokens:
return 0.0
kw = set([k.lower() for k in keywords or []])
kw = set([k.lower() for k in (keywords or [])])
inter = len(query_tokens & kw)
union = max(len(query_tokens | kw), 1)
return inter * 4 + (inter / union) * 100 + weight * 3
def emit_rows(rows, max_chars):
def dot(a: List[float], b: List[float]) -> float:
return sum(x * y for x, y in zip(a, b))
def norm(a: List[float]) -> float:
return math.sqrt(sum(x * x for x in a))
def cosine_similarity(a: List[float], b: List[float]) -> float:
na = norm(a)
nb = norm(b)
if na == 0.0 or nb == 0.0:
return 0.0
return dot(a, b) / (na * nb)
def embed_text(base_url: str, api_key: str, model: str, text: str) -> List[float]:
url = base_url.rstrip("/") + "/embeddings"
payload = json.dumps({
"model": model,
"input": [text]
}).encode("utf-8")
req = urllib.request.Request(
url,
data=payload,
headers={
"Content-Type": "application/json",
"Authorization": f"Bearer {api_key}"
},
method="POST"
)
with urllib.request.urlopen(req, timeout=180) as resp:
data = json.loads(resp.read().decode("utf-8"))
items = data.get("data")
if not isinstance(items, list) or not items:
raise RuntimeError("Invalid embedding response")
vec = items[0].get("embedding")
if not isinstance(vec, list):
raise RuntimeError("Missing embedding vector")
return vec
def build_query_text(rewrite: Dict, issue: Dict) -> str:
parts = []
parts.extend(rewrite.get("search_queries", []))
parts.extend(rewrite.get("keywords", []))
parts.extend(rewrite.get("components", []))
parts.extend(rewrite.get("likely_paths", []))
parts.append(issue.get("title", ""))
parts.append(issue.get("body", ""))
return "\n".join([p for p in parts if isinstance(p, str) and p.strip()])
def emit_rows(rows: List[Dict], max_chars: int):
parts = []
cur = 0
seen = set()
for row in rows:
key = (row["path"], row["start_line"], row["end_line"])
if key in seen:
continue
seen.add(key)
block = f"--- FILE: {row['path']} (lines {row['start_line']}-{row['end_line']}) ---\n{row['text']}\n"
if cur + len(block) > max_chars:
break
parts.append(block)
cur += len(block)
print("\n".join(parts))
def retrieve_from_index(index_path, rewrite, issue, max_chars):
queries = " ".join(rewrite.get("search_queries", []) + rewrite.get("keywords", []) + rewrite.get("components", []) + rewrite.get("likely_paths", []))
q = set(tokenize(queries + "\n" + issue.get("title", "") + "\n" + issue.get("body", "")))
def retrieve_from_index(index_path: str, rewrite: Dict, issue: Dict, max_chars: int):
query_text = build_query_text(rewrite, issue)
q = set(tokenize(query_text))
rows = []
with open(index_path, "r", encoding="utf-8") as f:
@@ -48,19 +127,71 @@ def retrieve_from_index(index_path, rewrite, issue, max_chars):
rows.sort(key=lambda x: (-x["score"], x["path"], x["start_line"]))
emit_rows(rows[:20], max_chars)
def retrieve_from_repo(repo_root, rewrite, issue, max_chars):
def retrieve_from_index_with_embeddings(
index_path: str,
embedding_index_path: str,
rewrite: Dict,
issue: Dict,
max_chars: int,
embedding_api_key: str,
embedding_base_url: str,
embedding_model: str
):
query_text = build_query_text(rewrite, issue)
q_tokens = set(tokenize(query_text))
q_vec = embed_text(embedding_base_url, embedding_api_key, embedding_model, query_text)
base_rows: Dict[str, Dict] = {}
with open(index_path, "r", encoding="utf-8") as f:
for line in f:
if not line.strip():
continue
row = json.loads(line)
base_rows[row["id"]] = row
scored = []
with open(embedding_index_path, "r", encoding="utf-8") as f:
for line in f:
if not line.strip():
continue
item = json.loads(line)
rid = item.get("id")
if rid not in base_rows:
continue
base = base_rows[rid]
emb = item.get("embedding")
if not isinstance(emb, list):
continue
sim = cosine_similarity(q_vec, emb)
kw_score = score_keywords(q_tokens, base.get("keywords"), base.get("weight", 1))
fused = sim * 100 + kw_score
row = dict(base)
row["score"] = fused
scored.append(row)
scored.sort(key=lambda x: (-x["score"], x["path"], x["start_line"]))
emit_rows(scored[:20], max_chars)
def retrieve_from_repo(repo_root: str, rewrite: Dict, issue: Dict, max_chars: int):
root = pathlib.Path(repo_root)
query_terms = rewrite.get("keywords", []) + rewrite.get("components", []) + rewrite.get("likely_paths", [])
query_terms = [x for x in query_terms if isinstance(x, str) and x.strip()]
rows = []
for path in root.rglob("*"):
if not path.is_file():
continue
rel = path.relative_to(root).as_posix()
if any(part in rel for part in [".git/", "node_modules/", "dist/", "build/", "coverage/", ".next/", ".nuxt/"]):
continue
if path.suffix.lower() not in {".md",".mdx",".rst",".txt",".js",".jsx",".ts",".tsx",".py",".go",".java",".rs",".php",".rb",".yml",".yaml",".json",".toml",".sh",".vue"} and not path.name.startswith("README"):
if path.suffix.lower() not in TEXT_EXTS and not path.name.startswith("README") and not path.name.startswith("CHANGELOG"):
continue
try:
@@ -71,9 +202,10 @@ def retrieve_from_repo(repo_root, rewrite, issue, max_chars):
lower = text.lower()
score = 0
for t in query_terms:
if t.lower() in lower:
tl = t.lower()
if tl in lower:
score += 2
if t.lower() in rel.lower():
if tl in rel.lower():
score += 4
if score <= 0:
@@ -88,6 +220,7 @@ def retrieve_from_repo(repo_root, rewrite, issue, max_chars):
hit += 1
if hit <= 0:
continue
start = max(0, i - 18)
end = min(len(lines), i + 19)
block = "\n".join(f"{j+1:>5}: {lines[j]}" for j in range(start, end))
@@ -102,22 +235,46 @@ def retrieve_from_repo(repo_root, rewrite, issue, max_chars):
rows.sort(key=lambda x: (-x["score"], x["path"], x["start_line"]))
emit_rows(rows[:20], max_chars)
def main():
ap = argparse.ArgumentParser()
ap.add_argument("--index", default="")
ap.add_argument("--embedding-index", default="")
ap.add_argument("--repo-root", default="")
ap.add_argument("--rewrite", required=True)
ap.add_argument("--issue", required=True)
ap.add_argument("--max-chars", type=int, default=50000)
ap.add_argument("--embedding-api-key", default="")
ap.add_argument("--embedding-base-url", default="")
ap.add_argument("--embedding-model", default="")
args = ap.parse_args()
rewrite = load_json(args.rewrite)
issue = load_json(args.issue)
if args.index and args.embedding_index and args.embedding_api_key and args.embedding_base_url and args.embedding_model:
retrieve_from_index_with_embeddings(
args.index,
args.embedding_index,
rewrite,
issue,
args.max_chars,
args.embedding_api_key,
args.embedding_base_url,
args.embedding_model
)
return
if args.index:
retrieve_from_index(args.index, rewrite, issue, args.max_chars)
else:
return
if args.repo_root:
retrieve_from_repo(args.repo_root, rewrite, issue, args.max_chars)
return
print("")
if __name__ == "__main__":
main()