Update generators/qa.py

generators/qa.py

@@ -1,28 +1,50 @@
-import json
-from rag.retriever import retrieve, format_citations
-from rag.prompts import QA_SYS, QA_USER
-from irpr.deps import generate_chat
+# generators/qa.py
+from __future__ import annotations
+from typing import List, Tuple
+from irpr.deps import search, generate_chat
 
-def make_qa(query: str, n=30):
-    hits = retrieve(query, top_k=16)
-    links = format_citations(hits)
-    link_lines = "\n".join([f"[{i}] {u}" for i, u in links])
-    ctx = "\n".join([f"[chunk:{h['chunk_id']}] {h['text']}" for h in hits])
+SYS = "あなたは日本語のIR担当です。投資家からの想定質問と模範回答を、根拠に基づいて簡潔に作成します。"
 
-    messages = [
-        {"role": "system", "content": QA_SYS},
-        {"role": "user", "content": QA_USER.format(links=link_lines, contexts=ctx)}
-    ]
-    raw = generate_chat(messages, max_new_tokens=1800, temperature=0.2)
+TPL = """次の抜粋を根拠に、投資家向けの想定Q&Aを {n} 問作成してください。
+各問は "Q: ... / A: ..." の2行で、Aは2-4文以内。根拠があれば括弧で短く示してください。
 
+# 抜粋
+{context}
+"""
+
+def make_qa(query: str, n: int = 30) -> Tuple[List[dict], List[str]]:
+    hits = search(query, top_k=min(12, n))
+    links, ctx = [], []
+    for i, h in enumerate(hits, 1):
+        src = h.get("source_url") or ""
+        if src and src not in links: links.append(src)
+        ctx.append(f"[{i}] {h.get('title') or ''} {src}\n{h['text'][:1000]}")
+    context = "\n\n".join(ctx) if ctx else "（根拠なし）"
+
+    # LLM試行
+    qa_list: List[dict] = []
     try:
-        data = json.loads(raw)
+        out = generate_chat(
+            [{"role":"system","content":SYS},
+             {"role":"user","content":TPL.format(n=n, context=context)}],
+            max_new_tokens=1000
+        )
+        for line in out.splitlines():
+            line = line.strip()
+            if line.startswith("Q:"):
+                qa_list.append({"q": line[2:].strip(), "a": ""})
+            elif line.startswith("A:") and qa_list:
+                qa_list[-1]["a"] = line[2:].strip()
+        qa_list = [x for x in qa_list if x.get("q") and x.get("a")]
+        if qa_list:
+            return qa_list[:n], links
     except Exception:
-        data = []
-        for line in raw.splitlines():
-            if "Q:" in line and "A:" in line:
-                q = line.split("Q:", 1)[1].split("A:", 1)[0].strip()
-                a = line.split("A:", 1)[1].strip()
-                data.append({"q": q, "a": a, "sources": []})
+        pass
 
-    return data[:n], links
+    # フォールバック（抽出）
+    for h in hits[:n]:
+        qa_list.append({
+            "q": f"{(h.get('title') or '決算トピック')}のポイントは？",
+            "a": (h['text'][:240] + "…").replace("\n"," "),
+        })
+    return qa_list[:n], links