171 lines
5.3 KiB
Python
171 lines
5.3 KiB
Python
"""
|
|
Copyright 2024-2026 ReplyPier
|
|
|
|
Licensed under the Apache License, Version 2.0 (the "License");
|
|
you may not use this file except in compliance with the License.
|
|
You may obtain a copy of the License at
|
|
|
|
http://www.apache.org/licenses/LICENSE-2.0
|
|
|
|
Unless required by applicable law or agreed to in writing, software
|
|
distributed under the License is distributed on an "AS IS" BASIS,
|
|
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
See the License for the specific language governing permissions and
|
|
limitations under the License.
|
|
"""
|
|
|
|
import json
|
|
from unittest.mock import MagicMock
|
|
from uuid import uuid4
|
|
|
|
from sqlalchemy import text
|
|
from sqlalchemy.dialects import postgresql
|
|
|
|
from app.services import knowledge_vector_links as kvl
|
|
|
|
|
|
def _knowledge(source="https://example.com"):
|
|
k = MagicMock()
|
|
k.source = source
|
|
k.schema = "ai"
|
|
k.table_name = "d_00000000-0000-0000-0000-000000000000"
|
|
k.organization_id = uuid4()
|
|
return k
|
|
|
|
|
|
def _db_returning(filters):
|
|
"""A db whose SELECT yields one row with the given filters."""
|
|
db = MagicMock()
|
|
row = MagicMock()
|
|
row.filters = filters
|
|
db.execute.return_value.first.return_value = row
|
|
return db
|
|
|
|
|
|
def test_update_sql_binds_every_parameter():
|
|
"""The regression: ':param::jsonb' is not parsed as a bind parameter.
|
|
|
|
SQLAlchemy's text() bind regex has a negative lookahead on ':', so a
|
|
parameter immediately followed by the '::' cast operator is left in the
|
|
string verbatim and reaches psycopg as a literal ':new_filters', which
|
|
Postgres rejects with `syntax error at or near ":"`. CAST(... AS jsonb)
|
|
avoids the adjacency.
|
|
"""
|
|
sql = kvl._UPDATE_SQL.format(schema="ai", table="d_x")
|
|
compiled = text(sql).compile(dialect=postgresql.dialect())
|
|
|
|
assert set(compiled.params) == {"new_filters", "agent_ids", "has_agents", "source"}
|
|
# Nothing left for the driver to choke on.
|
|
assert ":new_filters" not in str(compiled)
|
|
assert ":agent_ids" not in str(compiled)
|
|
|
|
|
|
def test_update_sql_uses_no_double_colon_cast():
|
|
"""Guard the specific construct, so it cannot be reintroduced."""
|
|
assert "::" not in kvl._UPDATE_SQL
|
|
|
|
|
|
def test_add_agent_preserves_existing_agents():
|
|
"""Linking a second agent must not evict the first.
|
|
|
|
The old code did `agent_ids = [new]` whenever the id was already present or
|
|
the value was not a list, which wiped every other agent from the source.
|
|
"""
|
|
db = _db_returning({"agent_id": ["aaa", "bbb"]})
|
|
new_agent = uuid4()
|
|
|
|
kvl.add_agent(db, _knowledge(), new_agent)
|
|
|
|
params = db.execute.call_args.args[1]
|
|
assert json.loads(params["agent_ids"]) == sorted(["aaa", "bbb", str(new_agent)])
|
|
|
|
|
|
def test_add_agent_is_idempotent():
|
|
"""Re-adding an already-present agent is a no-op, not a reset."""
|
|
existing = uuid4()
|
|
db = _db_returning({"agent_id": [str(existing), "other"]})
|
|
|
|
kvl.add_agent(db, _knowledge(), existing)
|
|
|
|
params = db.execute.call_args.args[1]
|
|
assert json.loads(params["agent_ids"]) == sorted([str(existing), "other"])
|
|
|
|
|
|
def test_remove_agent_keeps_the_others():
|
|
keep, drop = str(uuid4()), uuid4()
|
|
db = _db_returning({"agent_id": [keep, str(drop)]})
|
|
|
|
kvl.remove_agent(db, _knowledge(), drop)
|
|
|
|
params = db.execute.call_args.args[1]
|
|
assert json.loads(params["agent_ids"]) == [keep]
|
|
assert params["has_agents"] is True
|
|
|
|
|
|
def test_remove_last_agent_yields_empty_list():
|
|
only = uuid4()
|
|
db = _db_returning({"agent_id": [str(only)]})
|
|
|
|
kvl.remove_agent(db, _knowledge(), only)
|
|
|
|
params = db.execute.call_args.args[1]
|
|
assert json.loads(params["agent_ids"]) == []
|
|
assert params["has_agents"] is False
|
|
|
|
|
|
def test_filters_payload_carries_source_and_org():
|
|
db = _db_returning({"agent_id": []})
|
|
knowledge = _knowledge("https://docs.example.com")
|
|
agent = uuid4()
|
|
|
|
kvl.add_agent(db, knowledge, agent)
|
|
|
|
filters = json.loads(db.execute.call_args.args[1]["new_filters"])
|
|
assert filters["name"] == "https://docs.example.com"
|
|
assert filters["org_id"] == str(knowledge.organization_id)
|
|
assert filters["agent_id"] == [str(agent)]
|
|
|
|
|
|
def test_source_is_bound_never_interpolated():
|
|
"""A quote in the title must not be able to break out of the statement."""
|
|
db = _db_returning({"agent_id": []})
|
|
knowledge = _knowledge("it's \"quoted\"; DROP TABLE x;--")
|
|
|
|
kvl.add_agent(db, knowledge, uuid4())
|
|
|
|
sql, params = db.execute.call_args.args
|
|
assert params["source"] == "it's \"quoted\"; DROP TABLE x;--"
|
|
assert "DROP TABLE" not in str(sql)
|
|
|
|
|
|
def test_no_vector_rows_is_a_no_op():
|
|
"""A source whose chunks are not written yet must not be updated."""
|
|
db = MagicMock()
|
|
db.execute.return_value.first.return_value = None
|
|
|
|
kvl.add_agent(db, _knowledge(), uuid4())
|
|
|
|
# Only the SELECT ran; no UPDATE was issued.
|
|
assert db.execute.call_count == 1
|
|
|
|
|
|
def test_missing_table_metadata_is_a_no_op():
|
|
db = MagicMock()
|
|
knowledge = _knowledge()
|
|
knowledge.table_name = None
|
|
|
|
kvl.add_agent(db, knowledge, uuid4())
|
|
|
|
db.execute.assert_not_called()
|
|
|
|
|
|
def test_non_list_agent_filter_is_normalised():
|
|
"""Corrupt/legacy JSON must not propagate a bad type into the rewrite."""
|
|
db = _db_returning({"agent_id": "not-a-list"})
|
|
agent = uuid4()
|
|
|
|
kvl.add_agent(db, _knowledge(), agent)
|
|
|
|
params = db.execute.call_args.args[1]
|
|
assert json.loads(params["agent_ids"]) == [str(agent)]
|