206 lines
8.2 KiB
Python
206 lines
8.2 KiB
Python
from __future__ import annotations
|
|
|
|
import json
|
|
import threading
|
|
import sqlite3
|
|
import sys
|
|
from pathlib import Path
|
|
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
|
|
|
|
|
|
ROOT = Path(__file__).resolve().parents[2]
|
|
sys.path.insert(0, str(ROOT / "scripts"))
|
|
|
|
from build_1c_rag_index import main as _unused_lexical_main # noqa: F401,E402
|
|
from build_1c_rag_vector_index import build_vector_index # noqa: E402
|
|
from check_1c_rag_vector_freshness import check_vector_freshness # noqa: E402
|
|
from common import build_lexical_index, write_json # noqa: E402
|
|
from search_1c_rag_hybrid import hybrid_search # noqa: E402
|
|
from search_1c_rag_vector import search_vector_index # noqa: E402
|
|
|
|
|
|
class EmbeddingHandler(BaseHTTPRequestHandler):
|
|
calls: list[dict] = []
|
|
|
|
def do_POST(self) -> None: # noqa: N802
|
|
length = int(self.headers.get("Content-Length") or "0")
|
|
payload = json.loads(self.rfile.read(length).decode("utf-8"))
|
|
EmbeddingHandler.calls.append({"path": self.path, "payload": payload, "authorization": self.headers.get("Authorization")})
|
|
inputs = payload.get("input") if isinstance(payload.get("input"), list) else [payload.get("input")]
|
|
data = []
|
|
for index, text in enumerate(inputs):
|
|
text_value = str(text or "").lower()
|
|
if "номенклатура" in text_value or "реквизит" in text_value:
|
|
vector = [1.0, 0.0, 0.0]
|
|
elif "передзаписью" in text_value:
|
|
vector = [0.0, 1.0, 0.0]
|
|
else:
|
|
vector = [0.0, 0.0, 1.0]
|
|
data.append({"object": "embedding", "index": index, "embedding": vector})
|
|
body = json.dumps({"object": "list", "data": data}).encode("utf-8")
|
|
self.send_response(200)
|
|
self.send_header("Content-Type", "application/json")
|
|
self.send_header("Content-Length", str(len(body)))
|
|
self.end_headers()
|
|
self.wfile.write(body)
|
|
|
|
def log_message(self, format: str, *args: object) -> None:
|
|
return
|
|
|
|
|
|
def run_embedding_server() -> tuple[ThreadingHTTPServer, str]:
|
|
EmbeddingHandler.calls = []
|
|
server = ThreadingHTTPServer(("127.0.0.1", 0), EmbeddingHandler)
|
|
thread = threading.Thread(target=server.serve_forever, daemon=True)
|
|
thread.start()
|
|
return server, f"http://127.0.0.1:{server.server_port}"
|
|
|
|
|
|
def write_jsonl(path: Path, records: list[dict]) -> None:
|
|
path.parent.mkdir(parents=True, exist_ok=True)
|
|
path.write_text("".join(json.dumps(record, ensure_ascii=False) + "\n" for record in records), encoding="utf-8")
|
|
|
|
|
|
def sample_records() -> list[dict]:
|
|
return [
|
|
{
|
|
"id": "doc-bsl-0",
|
|
"document_id": "doc-bsl",
|
|
"source_path": "bsl.md",
|
|
"title": "ПередЗаписью",
|
|
"chunk_index": 0,
|
|
"content": "Процедура ПередЗаписью проверяет заполнение Наименование и ставит Отказ.",
|
|
"metadata": {"source_type": "bsl", "heading": "ПередЗаписью"},
|
|
},
|
|
{
|
|
"id": "doc-meta-0",
|
|
"document_id": "doc-meta",
|
|
"source_path": "metadata.md",
|
|
"title": "Реквизиты справочника",
|
|
"chunk_index": 0,
|
|
"content": "Справочник Номенклатура содержит реквизиты Артикул и ВидНоменклатуры.",
|
|
"metadata": {"source_type": "metadata", "heading": "Реквизиты"},
|
|
},
|
|
]
|
|
|
|
|
|
def test_build_and_search_vector_index(tmp_path: Path) -> None:
|
|
corpus = tmp_path / "rag_corpus.jsonl"
|
|
index = tmp_path / "rag_vector_index.sqlite"
|
|
write_jsonl(corpus, sample_records())
|
|
|
|
result = build_vector_index(corpus, index, dimensions=64, embedding_model="local-hashing-v1")
|
|
search = search_vector_index(index, "реквизиты номенклатура", limit=2, source_types=["metadata"], corpus_path=corpus)
|
|
|
|
assert result["status"] == "ok"
|
|
assert result["doc_count"] == 2
|
|
assert index.exists()
|
|
assert search["status"] == "ok"
|
|
assert search["freshness"]["status"] == "fresh"
|
|
assert search["meta"]["embedding_model"] == "local-hashing-v1"
|
|
assert search["results"]
|
|
assert search["results"][0]["document"]["source_path"] == "metadata.md"
|
|
|
|
|
|
def test_vector_index_reports_stale_corpus(tmp_path: Path) -> None:
|
|
corpus = tmp_path / "rag_corpus.jsonl"
|
|
index = tmp_path / "rag_vector_index.sqlite"
|
|
records = sample_records()
|
|
write_jsonl(corpus, records)
|
|
build_vector_index(corpus, index, dimensions=64, embedding_model="local-hashing-v1")
|
|
records.append(
|
|
{
|
|
"id": "doc-query-0",
|
|
"document_id": "doc-query",
|
|
"source_path": "query.md",
|
|
"title": "Запрос",
|
|
"chunk_index": 0,
|
|
"content": "ВЫБРАТЬ первые записи из регистра.",
|
|
"metadata": {"source_type": "query"},
|
|
}
|
|
)
|
|
write_jsonl(corpus, records)
|
|
|
|
search = search_vector_index(index, "запрос выбрать", limit=2, corpus_path=corpus)
|
|
|
|
assert search["freshness"]["status"] == "stale"
|
|
|
|
|
|
def test_check_vector_freshness_reports_missing_and_fresh(tmp_path: Path) -> None:
|
|
corpus = tmp_path / "rag_corpus.jsonl"
|
|
index = tmp_path / "rag_vector_index.sqlite"
|
|
write_jsonl(corpus, sample_records())
|
|
|
|
missing = check_vector_freshness(index, corpus)
|
|
build_vector_index(corpus, index, dimensions=64, embedding_model="local-hashing-v1")
|
|
fresh = check_vector_freshness(index, corpus)
|
|
|
|
assert missing["status"] == "missing"
|
|
assert fresh["status"] == "fresh"
|
|
assert fresh["embedding_model"] == "local-hashing-v1"
|
|
assert fresh["doc_count"] == 2
|
|
|
|
|
|
def test_hybrid_search_combines_lexical_and_vector(tmp_path: Path) -> None:
|
|
corpus = tmp_path / "rag_corpus.jsonl"
|
|
lexical_index = tmp_path / "rag_index.json"
|
|
vector_index = tmp_path / "rag_vector_index.sqlite"
|
|
records = sample_records()
|
|
write_jsonl(corpus, records)
|
|
write_json(lexical_index, build_lexical_index(records))
|
|
build_vector_index(corpus, vector_index, dimensions=64, embedding_model="local-hashing-v1")
|
|
|
|
result = hybrid_search(
|
|
"проверка перед записью",
|
|
lexical_index_path=lexical_index,
|
|
vector_index_path=vector_index,
|
|
corpus_path=corpus,
|
|
limit=2,
|
|
candidate_limit=10,
|
|
source_types=None,
|
|
)
|
|
|
|
assert result["status"] == "ok"
|
|
assert result["vector_freshness"]["status"] == "fresh"
|
|
assert result["results"]
|
|
assert "lexical" in result["results"][0]["channels"] or "vector" in result["results"][0]["channels"]
|
|
|
|
|
|
def test_openai_compatible_embedding_provider_builds_and_searches(tmp_path: Path, monkeypatch) -> None:
|
|
corpus = tmp_path / "rag_corpus.jsonl"
|
|
index = tmp_path / "rag_vector_index.sqlite"
|
|
write_jsonl(corpus, sample_records())
|
|
server, base_url = run_embedding_server()
|
|
monkeypatch.setenv("TEST_EMBEDDING_KEY", "secret-test-key")
|
|
try:
|
|
result = build_vector_index(
|
|
corpus,
|
|
index,
|
|
dimensions=3,
|
|
embedding_model="test-embedding-model",
|
|
embedding_provider="openai-compatible",
|
|
embedding_base_url=base_url,
|
|
embedding_api_key_env="TEST_EMBEDDING_KEY",
|
|
batch_size=2,
|
|
)
|
|
search = search_vector_index(
|
|
index,
|
|
"реквизиты номенклатура",
|
|
limit=2,
|
|
corpus_path=corpus,
|
|
embedding_api_key_env="TEST_EMBEDDING_KEY",
|
|
)
|
|
finally:
|
|
server.shutdown()
|
|
server.server_close()
|
|
|
|
assert result["embedding_provider"] == "openai-compatible"
|
|
assert result["embedding_dimensions"] == 3
|
|
assert search["status"] == "ok"
|
|
assert search["meta"]["embedding_provider"] == "openai-compatible"
|
|
assert search["meta"].get("embedding_api_key") is None
|
|
assert search["results"][0]["document"]["source_path"] == "metadata.md"
|
|
assert EmbeddingHandler.calls
|
|
assert all(call["payload"]["model"] == "test-embedding-model" for call in EmbeddingHandler.calls)
|
|
assert any(call["authorization"] == "Bearer secret-test-key" for call in EmbeddingHandler.calls)
|