103 lines
3.0 KiB
Python
103 lines
3.0 KiB
Python
from __future__ import annotations
|
||
|
||
from pathlib import Path
|
||
from typing import Any
|
||
|
||
from common import ROOT, read_yaml_mapping
|
||
|
||
|
||
DEFAULT_PROFILES = ROOT / "plugins" / "1c" / "rag" / "profiles.yaml"
|
||
DEFAULT_PROFILE = "auto"
|
||
|
||
PROFILE_KEYWORDS = {
|
||
"metadata": [
|
||
"метадан",
|
||
"реквизит",
|
||
"табличн",
|
||
"справочник",
|
||
"документ",
|
||
"регистр",
|
||
"форма",
|
||
"объект",
|
||
"конфигурац",
|
||
],
|
||
"bsl": [
|
||
"bsl",
|
||
"бсл",
|
||
"код",
|
||
"процедур",
|
||
"функц",
|
||
"модул",
|
||
"ошибк",
|
||
"переменн",
|
||
"конецесли",
|
||
"передзаписью",
|
||
],
|
||
"query": [
|
||
"запрос",
|
||
"выбрать",
|
||
"срез",
|
||
"остатк",
|
||
"read-only",
|
||
"readonly",
|
||
"sql",
|
||
],
|
||
"safe-change": [
|
||
"измен",
|
||
"переимен",
|
||
"production",
|
||
"прод",
|
||
"резерв",
|
||
"backup",
|
||
"соглас",
|
||
"безопас",
|
||
"удалить",
|
||
],
|
||
}
|
||
|
||
|
||
def load_rag_profiles(path: Path = DEFAULT_PROFILES) -> dict[str, dict[str, Any]]:
|
||
data = read_yaml_mapping(path)
|
||
profiles = data.get("profiles")
|
||
if not isinstance(profiles, dict):
|
||
raise ValueError(f"{path} must contain a profiles mapping")
|
||
return {str(name): dict(value or {}) for name, value in profiles.items()}
|
||
|
||
|
||
def get_rag_profile(name: str | None, path: Path = DEFAULT_PROFILES) -> dict[str, Any]:
|
||
profiles = load_rag_profiles(path)
|
||
profile_name = name or DEFAULT_PROFILE
|
||
if profile_name not in profiles:
|
||
available = ", ".join(sorted(profiles))
|
||
raise ValueError(f"Unknown RAG profile `{profile_name}`. Available: {available}")
|
||
profile = profiles[profile_name]
|
||
profile["id"] = profile_name
|
||
profile.setdefault("source_types", [])
|
||
profile.setdefault("limit", 4)
|
||
profile.setdefault("candidate_limit", 40)
|
||
profile.setdefault("min_score", 0.0)
|
||
profile.setdefault("max_context_chars", 12000)
|
||
profile.setdefault("dedupe_by_document", False)
|
||
return profile
|
||
|
||
|
||
def detect_rag_profile(question: str) -> str:
|
||
text = question.lower().replace("ё", "е")
|
||
scores = {
|
||
profile: sum(1 for keyword in keywords if keyword in text)
|
||
for profile, keywords in PROFILE_KEYWORDS.items()
|
||
}
|
||
best_profile, best_score = max(scores.items(), key=lambda item: item[1])
|
||
return best_profile if best_score > 0 else "general"
|
||
|
||
|
||
def resolve_rag_profile(name: str | None, question: str, path: Path = DEFAULT_PROFILES) -> dict[str, Any]:
|
||
profile_name = name or DEFAULT_PROFILE
|
||
if profile_name == "auto":
|
||
profile_name = detect_rag_profile(question)
|
||
return get_rag_profile(profile_name, path)
|
||
|
||
|
||
def list_rag_profiles(path: Path = DEFAULT_PROFILES) -> list[dict[str, Any]]:
|
||
return [get_rag_profile(name, path) for name in sorted(load_rag_profiles(path))]
|