115 lines
4.7 KiB
Python
115 lines
4.7 KiB
Python
from __future__ import annotations
|
|
|
|
import argparse
|
|
import sys
|
|
import urllib.error
|
|
from pathlib import Path
|
|
|
|
from common import call_chat_completion, read_json, search_lexical_index
|
|
from rag_profiles import resolve_rag_profile
|
|
|
|
ROOT = Path(__file__).resolve().parents[1]
|
|
DEFAULT_INDEX = ROOT / "plugins" / "1c" / "datasets" / "prepared" / "rag_index.json"
|
|
DEFAULT_SYSTEM_PROMPT = ROOT / "plugins" / "1c" / "prompts" / "system.md"
|
|
DEFAULT_RAG_PROMPT = ROOT / "plugins" / "1c" / "prompts" / "rag-answer.md"
|
|
|
|
|
|
def format_context(results: list[dict], *, max_chars: int = 12000) -> str:
|
|
if not results:
|
|
return "Контекст не найден."
|
|
|
|
blocks = []
|
|
used_chars = 0
|
|
for position, result in enumerate(results, start=1):
|
|
document = result["document"]
|
|
source = document.get("source_path") or "unknown"
|
|
chunk = document.get("chunk_index")
|
|
title = document.get("title") or "unknown"
|
|
content = (document.get("content") or "").strip()
|
|
header = f"[{position}] source={source} title={title} chunk={chunk} score={result['score']:.4f}"
|
|
remaining = max_chars - used_chars - len(header) - 2
|
|
if remaining <= 0:
|
|
break
|
|
if len(content) > remaining:
|
|
content = content[: max(0, remaining - 3)].rstrip() + "..."
|
|
block = "\n".join([header, content])
|
|
blocks.append(block)
|
|
used_chars += len(block) + 2
|
|
return "\n\n".join(blocks)
|
|
|
|
|
|
def render_prompt(template_path: Path, context: str, question: str) -> str:
|
|
template = template_path.read_text(encoding="utf-8")
|
|
return template.replace("{{context}}", context).replace("{{question}}", question)
|
|
|
|
|
|
def main() -> int:
|
|
parser = argparse.ArgumentParser(description="Ask the 1C RAG assistant.")
|
|
parser.add_argument("question")
|
|
parser.add_argument("--index", type=Path, default=DEFAULT_INDEX)
|
|
parser.add_argument("--profile", default="auto")
|
|
parser.add_argument("--limit", type=int)
|
|
parser.add_argument("--candidate-limit", type=int)
|
|
parser.add_argument("--dedupe-by-document", action="store_true")
|
|
parser.add_argument("--min-score", type=float)
|
|
parser.add_argument("--source-type", action="append", dest="source_types")
|
|
parser.add_argument("--platform-version")
|
|
parser.add_argument("--platform-doc-id")
|
|
parser.add_argument("--max-context-chars", type=int)
|
|
parser.add_argument("--base-url", help="OpenAI-compatible endpoint, for example http://docker-gpu.cin.su:8000")
|
|
parser.add_argument("--model", default="qwen3-4b-instruct")
|
|
parser.add_argument("--print-prompt", action="store_true", help="Print assembled prompt instead of calling a model.")
|
|
parser.add_argument("--system-prompt", type=Path, default=DEFAULT_SYSTEM_PROMPT)
|
|
parser.add_argument("--rag-prompt", type=Path, default=DEFAULT_RAG_PROMPT)
|
|
args = parser.parse_args()
|
|
|
|
index = read_json(args.index)
|
|
profile = resolve_rag_profile(args.profile, args.question)
|
|
source_types = args.source_types if args.source_types is not None else profile["source_types"]
|
|
results = search_lexical_index(
|
|
index,
|
|
args.question,
|
|
limit=int(args.limit or profile["limit"]),
|
|
candidate_limit=int(args.candidate_limit or profile["candidate_limit"]),
|
|
dedupe_by_document=args.dedupe_by_document or bool(profile["dedupe_by_document"]),
|
|
min_score=float(args.min_score if args.min_score is not None else profile["min_score"]),
|
|
source_types=source_types,
|
|
metadata_filters={
|
|
"platform_version": args.platform_version or "",
|
|
"platform_doc_id": args.platform_doc_id or "",
|
|
},
|
|
)
|
|
context = format_context(results, max_chars=int(args.max_context_chars or profile["max_context_chars"]))
|
|
rag_prompt = render_prompt(args.rag_prompt, context=context, question=args.question)
|
|
|
|
if args.print_prompt or not args.base_url:
|
|
print(rag_prompt)
|
|
if not args.base_url and not args.print_prompt:
|
|
print(
|
|
"\nNo --base-url provided, so only the prompt was rendered.",
|
|
file=sys.stderr,
|
|
)
|
|
return 0
|
|
|
|
system_prompt = args.system_prompt.read_text(encoding="utf-8")
|
|
try:
|
|
answer = call_chat_completion(
|
|
base_url=args.base_url,
|
|
model=args.model,
|
|
messages=[
|
|
{"role": "system", "content": system_prompt},
|
|
{"role": "user", "content": rag_prompt},
|
|
],
|
|
max_tokens=1200,
|
|
)
|
|
except (urllib.error.URLError, ValueError) as exc:
|
|
print(f"Chat request failed: {exc}", file=sys.stderr)
|
|
return 1
|
|
|
|
print(answer)
|
|
return 0
|
|
|
|
|
|
if __name__ == "__main__":
|
|
raise SystemExit(main())
|