mirror of https://github.com/langgenius/dify.git
210 lines
7.4 KiB
Python
210 lines
7.4 KiB
Python
import json
|
|
import os
|
|
from typing import Any
|
|
|
|
import holo_search_sdk as holo
|
|
import pytest
|
|
from _pytest.monkeypatch import MonkeyPatch
|
|
from psycopg import sql as psql
|
|
|
|
# Shared in-memory storage: {table_name: {doc_id: {"id", "text", "meta", "embedding"}}}
|
|
_mock_tables: dict[str, dict[str, dict[str, Any]]] = {}
|
|
|
|
|
|
class MockSearchQuery:
|
|
"""Mock query builder for search_vector and search_text results."""
|
|
|
|
def __init__(self, table_name: str, search_type: str):
|
|
self._table_name = table_name
|
|
self._search_type = search_type
|
|
self._limit_val = 10
|
|
self._filter_sql = None
|
|
|
|
def select(self, columns):
|
|
return self
|
|
|
|
def limit(self, n):
|
|
self._limit_val = n
|
|
return self
|
|
|
|
def where(self, filter_sql):
|
|
self._filter_sql = filter_sql
|
|
return self
|
|
|
|
def _apply_filter(self, row: dict[str, Any]) -> bool:
|
|
"""Apply the filter SQL to check if a row matches."""
|
|
if self._filter_sql is None:
|
|
return True
|
|
|
|
# Extract literals (the document IDs) from the filter SQL
|
|
# Filter format: meta->>'document_id' IN ('doc1', 'doc2')
|
|
literals = [v for t, v in _extract_identifiers_and_literals(self._filter_sql) if t == "literal"]
|
|
if not literals:
|
|
return True
|
|
|
|
# Get the document_id from the row's meta field
|
|
meta = row.get("meta", "{}")
|
|
if isinstance(meta, str):
|
|
meta = json.loads(meta)
|
|
doc_id = meta.get("document_id")
|
|
|
|
return doc_id in literals
|
|
|
|
def fetchall(self):
|
|
data = _mock_tables.get(self._table_name, {})
|
|
results = []
|
|
for row in list(data.values())[: self._limit_val]:
|
|
# Apply filter if present
|
|
if not self._apply_filter(row):
|
|
continue
|
|
|
|
if self._search_type == "vector":
|
|
# row format expected by _process_vector_results: (distance, id, text, meta)
|
|
results.append((0.1, row["id"], row["text"], row["meta"]))
|
|
else:
|
|
# row format expected by _process_full_text_results: (id, text, meta, embedding, score)
|
|
results.append((row["id"], row["text"], row["meta"], row.get("embedding", []), 0.9))
|
|
return results
|
|
|
|
|
|
class MockTable:
|
|
"""Mock table object returned by client.open_table()."""
|
|
|
|
def __init__(self, table_name: str):
|
|
self._table_name = table_name
|
|
|
|
def upsert_multi(self, index_column, values, column_names, update=True, update_columns=None):
|
|
if self._table_name not in _mock_tables:
|
|
_mock_tables[self._table_name] = {}
|
|
id_idx = column_names.index("id")
|
|
for row in values:
|
|
doc_id = row[id_idx]
|
|
_mock_tables[self._table_name][doc_id] = dict(zip(column_names, row))
|
|
|
|
def search_vector(self, vector, column, distance_method, output_name):
|
|
return MockSearchQuery(self._table_name, "vector")
|
|
|
|
def search_text(self, column, expression, return_score=False, return_score_name="score", return_all_columns=False):
|
|
return MockSearchQuery(self._table_name, "text")
|
|
|
|
def set_vector_index(
|
|
self, column, distance_method, base_quantization_type, max_degree, ef_construction, use_reorder
|
|
):
|
|
pass
|
|
|
|
def create_text_index(self, index_name, column, tokenizer):
|
|
pass
|
|
|
|
|
|
def _extract_sql_template(query) -> str:
|
|
"""Extract the SQL template string from a psycopg Composed object."""
|
|
if isinstance(query, psql.Composed):
|
|
for part in query:
|
|
if isinstance(part, psql.SQL):
|
|
return part._obj
|
|
if isinstance(query, psql.SQL):
|
|
return query._obj
|
|
return ""
|
|
|
|
|
|
def _extract_identifiers_and_literals(query) -> list[Any]:
|
|
"""Extract Identifier and Literal values from a psycopg Composed object."""
|
|
values: list[Any] = []
|
|
if isinstance(query, psql.Composed):
|
|
for part in query:
|
|
if isinstance(part, psql.Identifier):
|
|
values.append(("ident", part._obj[0] if part._obj else ""))
|
|
elif isinstance(part, psql.Literal):
|
|
values.append(("literal", part._obj))
|
|
elif isinstance(part, psql.Composed):
|
|
# Handles SQL(...).join(...) for IN clauses
|
|
for sub in part:
|
|
if isinstance(sub, psql.Literal):
|
|
values.append(("literal", sub._obj))
|
|
return values
|
|
|
|
|
|
class MockHologresClient:
|
|
"""Mock holo_search_sdk client that stores data in memory."""
|
|
|
|
def connect(self):
|
|
pass
|
|
|
|
def check_table_exist(self, table_name):
|
|
return table_name in _mock_tables
|
|
|
|
def open_table(self, table_name):
|
|
return MockTable(table_name)
|
|
|
|
def execute(self, query, fetch_result=False):
|
|
template = _extract_sql_template(query)
|
|
params = _extract_identifiers_and_literals(query)
|
|
|
|
if "CREATE TABLE" in template.upper():
|
|
# Extract table name from first identifier
|
|
table_name = next((v for t, v in params if t == "ident"), "unknown")
|
|
if table_name not in _mock_tables:
|
|
_mock_tables[table_name] = {}
|
|
return None
|
|
|
|
if "SELECT 1" in template:
|
|
# text_exists: SELECT 1 FROM {table} WHERE id = {id} LIMIT 1
|
|
table_name = next((v for t, v in params if t == "ident"), "")
|
|
doc_id = next((v for t, v in params if t == "literal"), "")
|
|
data = _mock_tables.get(table_name, {})
|
|
return [(1,)] if doc_id in data else []
|
|
|
|
if "SELECT id" in template:
|
|
# get_ids_by_metadata_field: SELECT id FROM {table} WHERE meta->>{key} = {value}
|
|
table_name = next((v for t, v in params if t == "ident"), "")
|
|
literals = [v for t, v in params if t == "literal"]
|
|
key = literals[0] if len(literals) > 0 else ""
|
|
value = literals[1] if len(literals) > 1 else ""
|
|
data = _mock_tables.get(table_name, {})
|
|
return [(doc_id,) for doc_id, row in data.items() if json.loads(row.get("meta", "{}")).get(key) == value]
|
|
|
|
if "DELETE" in template.upper():
|
|
table_name = next((v for t, v in params if t == "ident"), "")
|
|
if "id IN" in template:
|
|
# delete_by_ids
|
|
ids_to_delete = [v for t, v in params if t == "literal"]
|
|
for did in ids_to_delete:
|
|
_mock_tables.get(table_name, {}).pop(did, None)
|
|
elif "meta->>" in template:
|
|
# delete_by_metadata_field
|
|
literals = [v for t, v in params if t == "literal"]
|
|
key = literals[0] if len(literals) > 0 else ""
|
|
value = literals[1] if len(literals) > 1 else ""
|
|
data = _mock_tables.get(table_name, {})
|
|
to_remove = [
|
|
doc_id for doc_id, row in data.items() if json.loads(row.get("meta", "{}")).get(key) == value
|
|
]
|
|
for did in to_remove:
|
|
data.pop(did, None)
|
|
return None
|
|
|
|
return [] if fetch_result else None
|
|
|
|
def drop_table(self, table_name):
|
|
_mock_tables.pop(table_name, None)
|
|
|
|
|
|
def mock_connect(**kwargs):
|
|
"""Replacement for holo_search_sdk.connect() that returns a mock client."""
|
|
return MockHologresClient()
|
|
|
|
|
|
MOCK = os.getenv("MOCK_SWITCH", "false").lower() == "true"
|
|
|
|
|
|
@pytest.fixture
|
|
def setup_hologres_mock(monkeypatch: MonkeyPatch):
|
|
if MOCK:
|
|
monkeypatch.setattr(holo, "connect", mock_connect)
|
|
|
|
yield
|
|
|
|
if MOCK:
|
|
_mock_tables.clear()
|
|
monkeypatch.undo()
|