init
This commit is contained in:
@@ -0,0 +1,13 @@
|
||||
"""核心服务模块"""
|
||||
|
||||
from .llm_factory import create_chat_model
|
||||
from .prompt_manager import get_prompt_manager
|
||||
from .sql_prompt_manager import get_sql_prompt_manager
|
||||
from .template_matcher import get_template_matcher
|
||||
|
||||
__all__ = [
|
||||
"create_chat_model",
|
||||
"get_prompt_manager",
|
||||
"get_sql_prompt_manager",
|
||||
"get_template_matcher",
|
||||
]
|
||||
@@ -0,0 +1,16 @@
|
||||
from typing import Optional
|
||||
from langchain_openai import ChatOpenAI
|
||||
from config import Config, MAX_RETRIES, TIMEOUT
|
||||
|
||||
|
||||
def create_chat_model(model_section: Optional[str] = None) -> ChatOpenAI:
|
||||
"""创建 LLM 实例"""
|
||||
model_config = Config.get_model_config(model_section)
|
||||
return ChatOpenAI(
|
||||
model=model_config['model'],
|
||||
api_key=model_config['api_key'],
|
||||
base_url=model_config.get('base_url'),
|
||||
temperature=0.1,
|
||||
max_retries=MAX_RETRIES,
|
||||
timeout=TIMEOUT
|
||||
)
|
||||
@@ -0,0 +1,42 @@
|
||||
import os
|
||||
from typing import Any, Dict, Optional
|
||||
|
||||
import yaml
|
||||
|
||||
|
||||
class PromptManager:
|
||||
"""提示词配置管理器"""
|
||||
|
||||
def __init__(self, config_path: Optional[str] = None):
|
||||
root_dir = os.path.dirname(os.path.dirname(os.path.dirname(__file__)))
|
||||
self._config_path = config_path or os.path.join(root_dir, "config", "prompts.yaml")
|
||||
self._data: Dict[str, Any] = {}
|
||||
self.reload()
|
||||
|
||||
def reload(self) -> None:
|
||||
"""重新加载提示词配置"""
|
||||
with open(self._config_path, "r", encoding="utf-8") as f:
|
||||
self._data = yaml.safe_load(f) or {}
|
||||
|
||||
def get(self, group: str, name: str, default: str = "") -> str:
|
||||
"""获取指定提示词"""
|
||||
return str(self._data.get(group, {}).get(name, default))
|
||||
|
||||
def list_groups(self) -> list[str]:
|
||||
"""列出所有分组"""
|
||||
return list(self._data.keys())
|
||||
|
||||
def list_prompts(self, group: str) -> list[str]:
|
||||
"""列出分组内提示词"""
|
||||
return list(self._data.get(group, {}).keys())
|
||||
|
||||
|
||||
_GLOBAL_PROMPT_MANAGER: Optional[PromptManager] = None
|
||||
|
||||
|
||||
def get_prompt_manager(config_path: Optional[str] = None) -> PromptManager:
|
||||
"""获取全局 PromptManager（单例）"""
|
||||
global _GLOBAL_PROMPT_MANAGER
|
||||
if _GLOBAL_PROMPT_MANAGER is None:
|
||||
_GLOBAL_PROMPT_MANAGER = PromptManager(config_path=config_path)
|
||||
return _GLOBAL_PROMPT_MANAGER
|
||||
@@ -0,0 +1,238 @@
|
||||
import json
|
||||
import os
|
||||
from typing import Any, Dict, List, Optional
|
||||
|
||||
from config import Config
|
||||
from services.storage.cache import CacheBase, NoopCache, RedisCache
|
||||
|
||||
|
||||
class SqlPromptManager:
|
||||
"""按表名读取 SQL 提示词，Redis 主存储 + 本地文件回退"""
|
||||
|
||||
KEY_PREFIX = "sql_prompt"
|
||||
TABLE_LIST_KEY = "sql_prompt:table_list"
|
||||
SOURCE_REDIS = "redis"
|
||||
SOURCE_FILE = "file"
|
||||
|
||||
def __init__(self, base_dir: Optional[str] = None):
|
||||
root_dir = os.path.dirname(os.path.dirname(os.path.dirname(__file__)))
|
||||
self._fallback_dir = base_dir or os.path.join(root_dir, "config", "sql_gen_prompts")
|
||||
self._cache = self._init_cache()
|
||||
self._cache_ttl = self._get_cache_ttl()
|
||||
self._use_redis_primary = self._get_use_redis_primary()
|
||||
|
||||
@staticmethod
|
||||
def _first_line(value: Any) -> str:
|
||||
if value is None:
|
||||
return ""
|
||||
text = str(value)
|
||||
return text.splitlines()[0].strip() if text else ""
|
||||
|
||||
@classmethod
|
||||
def _to_bool(cls, value: Any, default: bool = False) -> bool:
|
||||
text = cls._first_line(value).lower()
|
||||
if not text:
|
||||
return default
|
||||
return text in ("1", "true", "yes", "y", "on")
|
||||
|
||||
@classmethod
|
||||
def _to_int(cls, value: Any, default: int = 0) -> int:
|
||||
text = cls._first_line(value)
|
||||
if not text:
|
||||
return default
|
||||
try:
|
||||
return int(text)
|
||||
except Exception:
|
||||
return default
|
||||
|
||||
@classmethod
|
||||
def _extract_key_from_multiline_values(cls, redis_cfg: Dict[str, Any], target_key: str) -> str:
|
||||
token = f"{target_key}="
|
||||
for raw in redis_cfg.values():
|
||||
text = str(raw or "")
|
||||
for line in text.splitlines()[1:]:
|
||||
cleaned = line.strip()
|
||||
normalized = cleaned.replace(" ", "")
|
||||
if normalized.lower().startswith(token.lower()):
|
||||
return cleaned.split("=", 1)[1].strip()
|
||||
return ""
|
||||
|
||||
@classmethod
|
||||
def _get_cfg_value(cls, redis_cfg: Dict[str, Any], key: str, default: Any = "") -> Any:
|
||||
if key in redis_cfg:
|
||||
return redis_cfg.get(key, default)
|
||||
recovered = cls._extract_key_from_multiline_values(redis_cfg, key)
|
||||
return recovered if recovered else default
|
||||
|
||||
@classmethod
|
||||
def _get_cache_ttl(cls) -> Optional[int]:
|
||||
redis_cfg = Config.get_section("redis")
|
||||
ttl = cls._to_int(cls._get_cfg_value(redis_cfg, "sql_prompt_ttl", 0), default=0)
|
||||
if ttl:
|
||||
return ttl
|
||||
return None
|
||||
|
||||
@classmethod
|
||||
def _get_use_redis_primary(cls) -> bool:
|
||||
redis_cfg = Config.get_section("redis")
|
||||
enabled = cls._to_bool(cls._get_cfg_value(redis_cfg, "enabled", "false"), default=False)
|
||||
primary = cls._to_bool(cls._get_cfg_value(redis_cfg, "sql_prompt_redis_primary", "false"), default=False)
|
||||
return enabled and primary
|
||||
|
||||
@classmethod
|
||||
def _init_cache(cls):
|
||||
redis_cfg = Config.get_section("redis")
|
||||
enabled = cls._to_bool(cls._get_cfg_value(redis_cfg, "enabled", "false"), default=False)
|
||||
if not enabled:
|
||||
return NoopCache()
|
||||
|
||||
url = cls._first_line(cls._get_cfg_value(redis_cfg, "url"))
|
||||
db = cls._to_int(cls._get_cfg_value(redis_cfg, "db", cls._get_cfg_value(redis_cfg, "database", 0)), default=0)
|
||||
if not url:
|
||||
host = cls._first_line(cls._get_cfg_value(redis_cfg, "host"))
|
||||
port = cls._first_line(cls._get_cfg_value(redis_cfg, "port", "6379")) or "6379"
|
||||
password = cls._first_line(cls._get_cfg_value(redis_cfg, "password", ""))
|
||||
username = cls._first_line(cls._get_cfg_value(redis_cfg, "username", ""))
|
||||
database = cls._first_line(cls._get_cfg_value(redis_cfg, "database", str(db))) or str(db)
|
||||
if host:
|
||||
from urllib.parse import quote_plus
|
||||
if username and password:
|
||||
auth = f"{quote_plus(username)}:{quote_plus(password)}@"
|
||||
elif password:
|
||||
auth = f":{quote_plus(password)}@"
|
||||
else:
|
||||
auth = ""
|
||||
url = f"redis://{auth}{host}:{port}/{database}"
|
||||
|
||||
if not url:
|
||||
return NoopCache()
|
||||
try:
|
||||
return RedisCache(url=url, db=db)
|
||||
except Exception:
|
||||
return NoopCache()
|
||||
|
||||
@staticmethod
|
||||
def _safe_filename(name: str) -> str:
|
||||
return name.replace("..", "").replace("/", "_").replace("\\", "_")
|
||||
|
||||
def _redis_key(self, table_name: str) -> str:
|
||||
return f"{self.KEY_PREFIX}:{table_name}"
|
||||
|
||||
def get_prompt(self, table_name: str) -> Optional[Dict[str, Any]]:
|
||||
"""读取指定表的提示词，优先 Redis，回退本地文件"""
|
||||
if not table_name:
|
||||
return None
|
||||
|
||||
if self._use_redis_primary:
|
||||
prompt = self._get_from_redis(table_name)
|
||||
if prompt:
|
||||
return prompt
|
||||
|
||||
prompt = self._get_from_file(table_name)
|
||||
return prompt
|
||||
|
||||
def get_prompt_with_source(self, table_name: str) -> tuple[Optional[Dict[str, Any]], str]:
|
||||
"""读取指定表的提示词，返回 (prompt, source) 元组"""
|
||||
if not table_name:
|
||||
return None, self.SOURCE_FILE
|
||||
|
||||
if self._use_redis_primary:
|
||||
prompt = self._get_from_redis(table_name)
|
||||
if prompt:
|
||||
return prompt, self.SOURCE_REDIS
|
||||
|
||||
prompt = self._get_from_file(table_name)
|
||||
source = self.SOURCE_FILE if prompt else self.SOURCE_FILE
|
||||
return prompt, source
|
||||
|
||||
def _get_from_redis(self, table_name: str) -> Optional[Dict[str, Any]]:
|
||||
key = self._redis_key(table_name)
|
||||
try:
|
||||
data = self._cache.get(key)
|
||||
if data:
|
||||
try:
|
||||
return json.loads(data)
|
||||
except Exception:
|
||||
pass
|
||||
except Exception:
|
||||
pass
|
||||
return None
|
||||
|
||||
def _get_from_file(self, table_name: str) -> Optional[Dict[str, Any]]:
|
||||
safe_name = self._safe_filename(table_name)
|
||||
filename = safe_name + ".json"
|
||||
path = os.path.join(self._fallback_dir, filename)
|
||||
if not os.path.exists(path):
|
||||
return None
|
||||
|
||||
with open(path, "r", encoding="utf-8") as f:
|
||||
prompt = json.load(f)
|
||||
return prompt
|
||||
|
||||
def save_prompt(self, table_name: str, prompt: Dict[str, Any]) -> bool:
|
||||
"""保存提示词到 Redis"""
|
||||
if not table_name or not prompt:
|
||||
return False
|
||||
|
||||
key = self._redis_key(table_name)
|
||||
try:
|
||||
self._cache.set(key, json.dumps(prompt, ensure_ascii=False), self._cache_ttl)
|
||||
return True
|
||||
except Exception:
|
||||
return False
|
||||
|
||||
def delete_prompt(self, table_name: str) -> bool:
|
||||
"""从 Redis 删除提示词"""
|
||||
if not table_name:
|
||||
return False
|
||||
|
||||
key = self._redis_key(table_name)
|
||||
try:
|
||||
self._cache.delete(key)
|
||||
return True
|
||||
except Exception:
|
||||
return False
|
||||
|
||||
def list_tables(self) -> List[str]:
|
||||
"""列出 Redis 中所有表名"""
|
||||
pattern = f"{self.KEY_PREFIX}:*"
|
||||
keys = self._cache.keys(pattern)
|
||||
tables = []
|
||||
for key in keys:
|
||||
if key == self.TABLE_LIST_KEY:
|
||||
continue
|
||||
parts = key.split(":", 1)
|
||||
if len(parts) == 2:
|
||||
tables.append(parts[1])
|
||||
return tables
|
||||
|
||||
def sync_from_files(self, tables: Optional[List[str]] = None) -> Dict[str, bool]:
|
||||
"""从本地文件同步到 Redis"""
|
||||
results: Dict[str, bool] = {}
|
||||
|
||||
if tables:
|
||||
files_to_sync = [f"{self._safe_filename(t)}.json" for t in tables]
|
||||
else:
|
||||
try:
|
||||
files_to_sync = [f for f in os.listdir(self._fallback_dir) if f.endswith(".json")]
|
||||
except Exception:
|
||||
return results
|
||||
|
||||
for filename in files_to_sync:
|
||||
table_name = filename[:-5]
|
||||
prompt = self._get_from_file(table_name)
|
||||
if prompt:
|
||||
results[table_name] = self.save_prompt(table_name, prompt)
|
||||
|
||||
return results
|
||||
|
||||
|
||||
_GLOBAL_SQL_PROMPT_MANAGER: Optional[SqlPromptManager] = None
|
||||
|
||||
|
||||
def get_sql_prompt_manager(base_dir: Optional[str] = None) -> SqlPromptManager:
|
||||
"""获取全局 SqlPromptManager（单例）"""
|
||||
global _GLOBAL_SQL_PROMPT_MANAGER
|
||||
if _GLOBAL_SQL_PROMPT_MANAGER is None:
|
||||
_GLOBAL_SQL_PROMPT_MANAGER = SqlPromptManager(base_dir=base_dir)
|
||||
return _GLOBAL_SQL_PROMPT_MANAGER
|
||||
@@ -0,0 +1,201 @@
|
||||
import json
|
||||
import re
|
||||
from pathlib import Path
|
||||
from typing import Any, Dict, Optional, Set
|
||||
|
||||
from config import Config
|
||||
from services.core.sql_prompt_manager import get_sql_prompt_manager
|
||||
from services.integrations.ragflow_client import RagflowClient, extract_table_name
|
||||
|
||||
|
||||
IDENTIFIER_RE = re.compile(r"[a-zA-Z_][a-zA-Z0-9_]*")
|
||||
EXPLICIT_FILTER_FIELD_RE = re.compile(r"([a-zA-Z_][a-zA-Z0-9_]*)\s*=")
|
||||
|
||||
|
||||
def _normalize_term(term: str) -> str:
|
||||
return str(term or "").strip().lower()
|
||||
|
||||
|
||||
class TemplateMatcher:
|
||||
"""模板匹配器：RAGFlow 检索"""
|
||||
|
||||
KEYWORD_MATCH_BONUS = 50
|
||||
|
||||
def __init__(self):
|
||||
self._ragflow = RagflowClient()
|
||||
self._sql_prompt_manager = get_sql_prompt_manager()
|
||||
cfg = Config.get_section("ragflow")
|
||||
self._dataset_id = (cfg.get("table_retrieval_dataset_id") or "").strip()
|
||||
self._top_k = int(cfg.get("retrieval_top_k", 3))
|
||||
self._non_empty_tables, self._table_keywords = self._load_table_config()
|
||||
self._table_terms_cache: Dict[str, Set[str]] = {}
|
||||
|
||||
def _load_table_config(self) -> tuple[Optional[Set[str]], Dict[str, Set[str]]]:
|
||||
"""从本地 tables.json 读取非空模板表集合和关键词映射。"""
|
||||
try:
|
||||
tables_path = Path(__file__).resolve().parents[2] / "config" / "table_retrieval_prompts" / "tables.json"
|
||||
with open(tables_path, "r", encoding="utf-8") as f:
|
||||
data = json.load(f)
|
||||
|
||||
if not isinstance(data, dict):
|
||||
return None, {}
|
||||
|
||||
non_empty: Set[str] = set()
|
||||
keywords_map: Dict[str, Set[str]] = {}
|
||||
|
||||
for table_name, templates in data.items():
|
||||
if not isinstance(table_name, str):
|
||||
continue
|
||||
if isinstance(templates, list) and len(templates) > 0:
|
||||
non_empty.add(table_name)
|
||||
keywords_map[table_name] = {_normalize_term(kw) for kw in templates if isinstance(kw, str)}
|
||||
|
||||
return non_empty, keywords_map
|
||||
except Exception:
|
||||
return None, {}
|
||||
|
||||
def _validate(self) -> None:
|
||||
if not self._dataset_id:
|
||||
raise RuntimeError("未配置 ragflow.table_retrieval_dataset_id，无法进行表名检索")
|
||||
|
||||
@staticmethod
|
||||
def _extract_query_terms(normalized_text: str) -> Set[str]:
|
||||
return {
|
||||
_normalize_term(match.group(0))
|
||||
for match in IDENTIFIER_RE.finditer(normalized_text or "")
|
||||
}
|
||||
|
||||
@staticmethod
|
||||
def _extract_explicit_filter_fields(normalized_text: str) -> Set[str]:
|
||||
return {
|
||||
_normalize_term(match.group(1))
|
||||
for match in EXPLICIT_FILTER_FIELD_RE.finditer(normalized_text or "")
|
||||
}
|
||||
|
||||
def _load_table_terms(self, table_name: str) -> Set[str]:
|
||||
cached = self._table_terms_cache.get(table_name)
|
||||
if cached is not None:
|
||||
return cached
|
||||
|
||||
prompt = self._sql_prompt_manager.get_prompt(table_name) or {}
|
||||
terms: Set[str] = set()
|
||||
|
||||
for field in (prompt.get("data_model_specification") or {}).get("fields_list") or []:
|
||||
if isinstance(field, str):
|
||||
terms.add(_normalize_term(field))
|
||||
|
||||
field_ref = prompt.get("field_mapping_reference") or {}
|
||||
self._collect_mapping_terms(field_ref, terms)
|
||||
|
||||
self._table_terms_cache[table_name] = terms
|
||||
return terms
|
||||
|
||||
def _collect_mapping_terms(self, node: Any, terms: Set[str]) -> None:
|
||||
if isinstance(node, dict):
|
||||
for key, value in node.items():
|
||||
if key == "alias" and isinstance(value, list):
|
||||
for alias in value:
|
||||
if isinstance(alias, str):
|
||||
terms.add(_normalize_term(alias))
|
||||
continue
|
||||
|
||||
if isinstance(value, dict):
|
||||
if "alias" in value or "type" in value:
|
||||
terms.add(_normalize_term(key))
|
||||
self._collect_mapping_terms(value, terms)
|
||||
elif isinstance(value, list):
|
||||
# 对字段列表直接入词，增强字段覆盖匹配
|
||||
if key.endswith("_fields") or key in {"fields_list", "list"}:
|
||||
for item in value:
|
||||
if isinstance(item, str):
|
||||
terms.add(_normalize_term(item))
|
||||
self._collect_mapping_terms(value, terms)
|
||||
elif isinstance(node, list):
|
||||
for item in node:
|
||||
self._collect_mapping_terms(item, terms)
|
||||
|
||||
def _rank_candidates(self, normalized_text: str, candidates: list[Dict[str, Any]]) -> list[Dict[str, Any]]:
|
||||
query_terms = self._extract_query_terms(normalized_text)
|
||||
explicit_fields = self._extract_explicit_filter_fields(normalized_text)
|
||||
|
||||
if not query_terms or not candidates:
|
||||
return candidates
|
||||
|
||||
ranked: list[Dict[str, Any]] = []
|
||||
for index, candidate in enumerate(candidates):
|
||||
table_name = candidate.get("table_name")
|
||||
if not table_name:
|
||||
continue
|
||||
|
||||
table_terms = self._load_table_terms(table_name)
|
||||
overlap = len(query_terms & table_terms)
|
||||
|
||||
missing_explicit_fields = len([field for field in explicit_fields if field not in table_terms])
|
||||
|
||||
keyword_bonus = 0
|
||||
table_keywords = self._table_keywords.get(table_name, set())
|
||||
if table_keywords:
|
||||
matched_keywords = query_terms & table_keywords
|
||||
keyword_bonus = len(matched_keywords) * self.KEYWORD_MATCH_BONUS
|
||||
|
||||
rank_score = overlap + keyword_bonus - (missing_explicit_fields * 5)
|
||||
|
||||
ranked.append({
|
||||
**candidate,
|
||||
"rank_score": rank_score,
|
||||
"rank_overlap": overlap,
|
||||
"rank_keyword_bonus": keyword_bonus,
|
||||
"rank_missing_explicit_fields": missing_explicit_fields,
|
||||
"rank_index": index,
|
||||
})
|
||||
|
||||
ranked.sort(key=lambda item: (item["rank_score"], item["rank_overlap"], -item["rank_index"]), reverse=True)
|
||||
return ranked
|
||||
|
||||
def match(self, normalized_text: str) -> Dict[str, Any]:
|
||||
"""返回匹配的表名与原始响应"""
|
||||
self._validate()
|
||||
try:
|
||||
response = self._ragflow.retrieve(normalized_text, top_k=self._top_k, dataset_id=self._dataset_id)
|
||||
except Exception as e:
|
||||
return {"table_name": None, "candidates": [], "raw": {"error": str(e)}}
|
||||
|
||||
candidates = []
|
||||
seen = set()
|
||||
data = response.get("data") if isinstance(response, dict) else None
|
||||
records = []
|
||||
if isinstance(data, list):
|
||||
records = data
|
||||
elif isinstance(data, dict):
|
||||
chunks = data.get("chunks")
|
||||
if isinstance(chunks, list):
|
||||
records = chunks
|
||||
|
||||
for item in records:
|
||||
table_name = extract_table_name(item)
|
||||
if self._non_empty_tables is not None and table_name not in self._non_empty_tables:
|
||||
continue
|
||||
if table_name and table_name not in seen:
|
||||
seen.add(table_name)
|
||||
candidates.append(
|
||||
{
|
||||
"table_name": table_name,
|
||||
"metadata": item.get("metadata") or {},
|
||||
"content": item.get("content") or item.get("text") or "",
|
||||
}
|
||||
)
|
||||
|
||||
candidates = self._rank_candidates(normalized_text, candidates)
|
||||
matched = candidates[0]["table_name"] if candidates else None
|
||||
return {"table_name": matched, "candidates": candidates, "raw": response}
|
||||
|
||||
|
||||
_GLOBAL_TEMPLATE_MATCHER: TemplateMatcher | None = None
|
||||
|
||||
|
||||
def get_template_matcher() -> TemplateMatcher:
|
||||
"""获取全局 TemplateMatcher（单例）"""
|
||||
global _GLOBAL_TEMPLATE_MATCHER
|
||||
if _GLOBAL_TEMPLATE_MATCHER is None:
|
||||
_GLOBAL_TEMPLATE_MATCHER = TemplateMatcher()
|
||||
return _GLOBAL_TEMPLATE_MATCHER
|
||||
Reference in New Issue
Block a user