← Back
Editing: llm.lua
--[[ LLM provider for neural feature fusion Collects text from the most relevant part and requests embeddings from an LLM API. Supports minimal OpenAI- and Ollama-compatible embedding endpoints. ]] -- local rspamd_http = require "rspamd_http" local rspamd_logger = require "rspamd_logger" local ucl = require "ucl" local neural_common = require "plugins/neural" local lua_cache = require "lua_cache" local llm_common = require "llm_common" local lua_mime = require "lua_mime" local N = "neural.llm" local function select_text(task, opts) return llm_common.build_llm_input(task, opts) end -- Detect primary language from the displayed text part local function detect_language(task) local part = lua_mime.get_displayed_text_part(task) if part then local lang = part:get_language() if lang and lang ~= '' then return lang end end return nil end local function compose_llm_settings(pcfg, language) local gpt_settings = rspamd_config:get_all_opt('gpt') or {} -- Provider identity is pcfg.type=='llm'; backend type is specified via one of these keys local llm_type = pcfg.llm_type or pcfg.api or pcfg.backend or gpt_settings.type or 'openai' local model = pcfg.model or gpt_settings.model local model_params = gpt_settings.model_parameters or {} local model_cfg = model and model_params[model] or {} local max_tokens = pcfg.max_tokens if not max_tokens then max_tokens = model_cfg.max_completion_tokens or model_cfg.max_tokens or gpt_settings.max_tokens end local timeout = pcfg.timeout or gpt_settings.timeout or 2.0 local url = pcfg.url local api_key = pcfg.api_key or gpt_settings.api_key -- Language-specific model/URL selection -- Config format: language_models = { en = { model = "...", url = "..." }, ru = { model = "..." }, ... } -- Or shorthand: language_models = { en = "model-name", ru = "model-name", ... } local language_models = pcfg.language_models if language and language_models then local lang_cfg = language_models[language] if lang_cfg then if type(lang_cfg) == 'string' then -- Shorthand: just model name model = lang_cfg elseif type(lang_cfg) == 'table' then -- Full config: { model = "...", url = "...", api_key = "..." } if lang_cfg.model then model = lang_cfg.model end if lang_cfg.url then url = lang_cfg.url end if lang_cfg.api_key then api_key = lang_cfg.api_key end end end end if not url then if llm_type == 'openai' then url = 'https://api.openai.com/v1/embeddings' elseif llm_type == 'ollama' then url = 'http://127.0.0.1:11434/api/embeddings' end end return { type = llm_type, model = model, max_tokens = max_tokens, timeout = timeout, url = url, api_key = api_key, cache_ttl = pcfg.cache_ttl or 86400, cache_prefix = pcfg.cache_prefix or 'neural_llm', cache_hash_len = pcfg.cache_hash_len or 32, cache_use_hashing = (pcfg.cache_use_hashing ~= false), -- Optional staged timeouts (inherit from global gpt if present) connect_timeout = pcfg.connect_timeout or gpt_settings.connect_timeout, ssl_timeout = pcfg.ssl_timeout or gpt_settings.ssl_timeout, write_timeout = pcfg.write_timeout or gpt_settings.write_timeout, read_timeout = pcfg.read_timeout or gpt_settings.read_timeout, reply_trim_mode = pcfg.reply_trim_mode or gpt_settings.reply_trim_mode, } end local function normalize_cache_key_input(input_string) if type(input_string) == 'userdata' then return input_string:str() end return tostring(input_string) end local function extract_embedding(llm_type, parsed) if llm_type == 'openai' then -- { data = [ { embedding = [...] } ] } if parsed and parsed.data and parsed.data[1] and parsed.data[1].embedding then return parsed.data[1].embedding end elseif llm_type == 'ollama' then -- { embedding = [...] } if parsed and parsed.embedding then return parsed.embedding end end return nil end neural_common.register_provider('llm', { collect_async = function(task, ctx, cont) local pcfg = ctx.config or {} -- Detect language from displayed text part for model/URL selection local language = detect_language(task) local llm = compose_llm_settings(pcfg, language) if not llm.model then rspamd_logger.debugm(N, task, 'llm provider missing model; skip') cont(nil) return end -- Do not run embeddings on infer if ANN is not loaded for this set/profile if ctx.phase == 'infer' then local set_or_profile = ctx.profile or ctx.set if not set_or_profile or not set_or_profile.ann then rspamd_logger.debugm(N, task, 'skip llm on infer: ANN not loaded for current settings') cont(nil) return end end local input_tbl = select_text(task, { reply_trim_mode = llm.reply_trim_mode }) if not input_tbl then rspamd_logger.debugm(N, task, 'llm provider has no content to embed; skip') cont(nil) return end -- Build request input string: subject first (more valuable for spam detection), -- then text content. Subject-first ensures it's always included even if text is truncated. local input_string if input_tbl.subject and input_tbl.subject ~= '' then input_string = "Subject: " .. input_tbl.subject .. "\n" .. (input_tbl.text or '') else input_string = input_tbl.text or '' end local input_key = normalize_cache_key_input(input_string) rspamd_logger.debugm(N, task, 'llm embedding request: model=%s url=%s lang=%s len=%s', tostring(llm.model), tostring(llm.url), tostring(language or 'unknown'), tostring(#input_key)) local body if llm.type == 'openai' then body = { model = llm.model, input = input_string } elseif llm.type == 'ollama' then body = { model = llm.model, prompt = input_string } else rspamd_logger.debugm(N, task, 'unsupported llm type: %s', llm.type) cont(nil) return end -- Redis cache: hash the final input string only local cache_ctx = lua_cache.create_cache_context(neural_common.redis_params, { cache_prefix = llm.cache_prefix, cache_ttl = llm.cache_ttl, cache_format = 'messagepack', cache_hash_len = llm.cache_hash_len, cache_use_hashing = llm.cache_use_hashing, }, N) -- Use raw key and allow cache module to hash/shorten it per context -- Include language in cache key for proper separation local key = string.format('%s:%s:%s:%s', llm.type, llm.model or 'model', language or 'unk', input_key) local function finish_with_vec(vec) if type(vec) == 'table' and #vec > 0 then local meta = { name = pcfg.name or 'llm', type = 'llm', dim = #vec, weight = ctx.weight or 1.0, model = llm.model, provider = llm.type, language = language, } rspamd_logger.debugm(N, task, 'llm embedding result: dim=%s lang=%s', #vec, language or 'unknown') cont(vec, meta) else rspamd_logger.debugm(N, task, 'llm embedding result: empty') cont(nil) end end local function http_cb(err, code, resp, _) if err then rspamd_logger.debugm(N, task, 'llm http error: %s', err) cont(nil) return end if code ~= 200 or not resp then rspamd_logger.debugm(N, task, 'llm bad http code: %s', code) cont(nil) return end local parser = ucl.parser() local ok, perr = parser:parse_string(resp) if not ok then rspamd_logger.debugm(N, task, 'llm cannot parse reply: %s', perr) cont(nil) return end local parsed = parser:get_object() local emb = extract_embedding(llm.type, parsed) if type(emb) == 'table' then lua_cache.cache_set(task, key, emb, cache_ctx) finish_with_vec(emb) else rspamd_logger.debugm(N, task, 'llm embedding parse: no embedding field') cont(nil) end end local function do_request_and_cache() local headers = { ['Content-Type'] = 'application/json' } if llm.type == 'openai' and llm.api_key then headers['Authorization'] = 'Bearer ' .. llm.api_key end local http_params = { url = llm.url, mime_type = 'application/json', timeout = llm.timeout, log_obj = task, headers = headers, body = ucl.to_format(body, 'json-compact', true), task = task, method = 'POST', use_gzip = true, keepalive = true, callback = http_cb, -- staged timeouts connect_timeout = llm.connect_timeout, ssl_timeout = llm.ssl_timeout, write_timeout = llm.write_timeout, read_timeout = llm.read_timeout, } rspamd_http.request(http_params) end -- Use async cache API lua_cache.cache_get(task, key, cache_ctx, llm.timeout or 2.0, function() -- Uncached path do_request_and_cache() end, function(_, err, data) if data and type(data) == 'table' then finish_with_vec(data) else do_request_and_cache() end end) end, })
Save File
Cancel