tiennm99/DocsGPT
1
0 Fork 0
mirror of https://github.com/tiennm99/DocsGPT.git synced 2026-10-14 12:16:20 +00:00
Code Issues Packages Projects Releases Wiki Activity

Pg 4 (#2390)

Browse Source 
* feat: postgres tests

* feat: mongo cutoff

* feat: mongo cutoff

* feat: adjust docs and compose files

* fix: mini code mongo removals

* fix: tests and k8s mongo stuff

* feat: test fixes

* fix: ruff

* fix: vale

* Potential fix for pull request finding 'CodeQL / Clear-text logging of sensitive information'

Co-authored-by: Copilot Autofix powered by AI <62310815+github-advanced-security[bot]@users.noreply.github.com>

* fix: mini suggestions

* vale lint fix 2

* fix: codeql columns thing

* fix: test mongo

* fix: tests coverage

* feat: better tests 4

* feat: more tests

* feat: decent coverage

* fix: ruff fixes

* fix: remove mongo mock

* feat: enhance workflow engine and API routes; add document retrieval and source handling

* feat: e2e tests

* fix: mcp, mongo and more

* fix: mini codeql warning

* fix: agent chunk view

* fix: mini issues

* fix: more pg fixes

* feat: postgres prep on start

* feat: qa tests

* fix: mini improvements

* fix: tests

---------

Co-authored-by: Copilot Autofix powered by AI <62310815+github-advanced-security[bot]@users.noreply.github.com>
Co-authored-by: Siddhant Rai <siddhant.rai.5686@gmail.com>
This commit is contained in:
AlexCopilot Autofix powered by AISiddhant Rai
 authored and GitHub committed 2026-04-18 13:13:57 +01:00
1 parent 7f6360b4ff
commit 81b6ee5daa
312 files changed
+763725 -25159
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

tests/api/answer/services/test_continuation_service_pg.py
+256
View File
@@ -0,0 +1,256 @@
"""Tests for application/api/answer/services/continuation_service.py using pg_conn."""
from contextlib import contextmanager
from unittest.mock import patch
from uuid import uuid4
@contextmanager
def _patch_db(conn):
@contextmanager
def _yield():
yield conn
with patch(
"application.api.answer.services.continuation_service.db_readonly",
_yield,
), patch(
"application.api.answer.services.continuation_service.db_session",
_yield,
):
yield
class TestMakeSerializable:
def test_uuid_becomes_string(self):
from application.api.answer.services.continuation_service import (
_make_serializable,
)
u = uuid4()
assert _make_serializable(u) == str(u)
def test_dict_keys_stringified(self):
from application.api.answer.services.continuation_service import (
_make_serializable,
)
got = _make_serializable({42: "a", "b": 2})
assert got == {"42": "a", "b": 2}
def test_list_elements_recursively_serialized(self):
from application.api.answer.services.continuation_service import (
_make_serializable,
)
u = uuid4()
got = _make_serializable([u, {"x": u}, 1])
assert got == [str(u), {"x": str(u)}, 1]
def test_bytes_decoded_to_string(self):
from application.api.answer.services.continuation_service import (
_make_serializable,
)
assert _make_serializable(b"hello") == "hello"
def test_bytes_invalid_utf8_replaced(self):
from application.api.answer.services.continuation_service import (
_make_serializable,
)
# Invalid UTF-8 byte sequence
got = _make_serializable(b"\xff\xfe")
assert isinstance(got, str)
def test_passes_through_primitives(self):
from application.api.answer.services.continuation_service import (
_make_serializable,
)
assert _make_serializable("hello") == "hello"
assert _make_serializable(42) == 42
assert _make_serializable(None) is None
assert _make_serializable(True) is True
class TestContinuationServiceSaveLoad:
def test_save_and_load_state(self, pg_conn):
from application.api.answer.services.continuation_service import (
ContinuationService,
)
from application.storage.db.repositories.conversations import (
ConversationsRepository,
)
user = "u-cont"
conv = ConversationsRepository(pg_conn).create(user, name="c")
conv_id = str(conv["id"])
service = ContinuationService()
with _patch_db(pg_conn):
service.save_state(
conversation_id=conv_id,
user=user,
messages=[{"role": "user", "content": "hi"}],
pending_tool_calls=[{"id": "call-1", "name": "search"}],
tools_dict={"search": {"arg": "value"}},
tool_schemas=[{"name": "search", "params": {}}],
agent_config={"model_id": "gpt-4"},
client_tools=[{"name": "client-tool"}],
)
loaded = service.load_state(conv_id, user)
assert loaded is not None
assert loaded["messages"] == [{"role": "user", "content": "hi"}]
def test_load_state_returns_none_when_not_found(self, pg_conn):
from application.api.answer.services.continuation_service import (
ContinuationService,
)
service = ContinuationService()
with _patch_db(pg_conn):
got = service.load_state(
"00000000-0000-0000-0000-000000000000", "u",
)
assert got is None
def test_save_state_no_client_tools(self, pg_conn):
from application.api.answer.services.continuation_service import (
ContinuationService,
)
from application.storage.db.repositories.conversations import (
ConversationsRepository,
)
user = "u-no-client"
conv = ConversationsRepository(pg_conn).create(user, name="c")
conv_id = str(conv["id"])
service = ContinuationService()
with _patch_db(pg_conn):
service.save_state(
conversation_id=conv_id,
user=user,
messages=[],
pending_tool_calls=[{"id": "c"}],
tools_dict={},
tool_schemas=[],
agent_config={},
client_tools=None,
)
loaded = service.load_state(conv_id, user)
assert loaded is not None
def test_delete_state_returns_true_when_deleted(self, pg_conn):
from application.api.answer.services.continuation_service import (
ContinuationService,
)
from application.storage.db.repositories.conversations import (
ConversationsRepository,
)
user = "u-del"
conv = ConversationsRepository(pg_conn).create(user, name="c")
conv_id = str(conv["id"])
service = ContinuationService()
with _patch_db(pg_conn):
service.save_state(
conversation_id=conv_id,
user=user,
messages=[],
pending_tool_calls=[{"id": "c"}],
tools_dict={},
tool_schemas=[],
agent_config={},
)
got = service.delete_state(conv_id, user)
assert got is True
def test_delete_state_returns_false_when_missing(self, pg_conn):
from application.api.answer.services.continuation_service import (
ContinuationService,
)
service = ContinuationService()
with _patch_db(pg_conn):
got = service.delete_state(
"00000000-0000-0000-0000-000000000000", "u",
)
assert got is False
class TestContinuationServiceLegacyIdHandling:
"""An unresolvable Mongo ObjectId must not reach ``CAST(:conv_id AS uuid)``
in ``PendingToolStateRepository`` — that cast raises and aborts the
enclosing transaction. This can happen when an OpenAI-compatible client
echoes back a ``chatcmpl-<legacy_objectid>`` id from before the PG
cutover and the conversation hasn't been backfilled."""
LEGACY_OBJECTID = "507f1f77bcf86cd799439011"
def test_load_state_unresolvable_legacy_id_returns_none(self, pg_conn):
from application.api.answer.services.continuation_service import (
ContinuationService,
)
service = ContinuationService()
with _patch_db(pg_conn):
got = service.load_state(self.LEGACY_OBJECTID, "u")
assert got is None
def test_delete_state_unresolvable_legacy_id_returns_false(self, pg_conn):
from application.api.answer.services.continuation_service import (
ContinuationService,
)
service = ContinuationService()
with _patch_db(pg_conn):
got = service.delete_state(self.LEGACY_OBJECTID, "u")
assert got is False
def test_save_state_unresolvable_legacy_id_raises(self, pg_conn):
import pytest
from application.api.answer.services.continuation_service import (
ContinuationService,
)
service = ContinuationService()
with _patch_db(pg_conn):
with pytest.raises(ValueError):
service.save_state(
conversation_id=self.LEGACY_OBJECTID,
user="u",
messages=[],
pending_tool_calls=[{"id": "c"}],
tools_dict={},
tool_schemas=[],
agent_config={},
)
def test_load_state_resolves_backfilled_legacy_id(self, pg_conn):
from application.api.answer.services.continuation_service import (
ContinuationService,
)
from application.storage.db.repositories.conversations import (
ConversationsRepository,
)
user = "legacy-cont"
ConversationsRepository(pg_conn).create(
user, name="c", legacy_mongo_id=self.LEGACY_OBJECTID,
)
service = ContinuationService()
with _patch_db(pg_conn):
# Save via the legacy id — should resolve to the PG UUID internally.
service.save_state(
conversation_id=self.LEGACY_OBJECTID,
user=user,
messages=[{"role": "user", "content": "hi"}],
pending_tool_calls=[{"id": "c"}],
tools_dict={},
tool_schemas=[],
agent_config={},
)
# And load via the legacy id — should round-trip.
loaded = service.load_state(self.LEGACY_OBJECTID, user)
assert loaded is not None
assert loaded["messages"] == [{"role": "user", "content": "hi"}]
tests/api/answer/services/test_conversation_service.py
+351 -218
View File
@@ -1,242 +1,375 @@
from unittest.mock import Mock
"""Tests for application/api/answer/services/conversation_service.py."""
from contextlib import contextmanager
from datetime import datetime, timezone
from unittest.mock import MagicMock, patch
import pytest
from bson import ObjectId
@pytest.mark.unit
@contextmanager
def _patch_db(conn):
@contextmanager
def _yield():
yield conn
with patch(
"application.api.answer.services.conversation_service.db_session",
_yield,
), patch(
"application.api.answer.services.conversation_service.db_readonly",
_yield,
):
yield
class TestConversationServiceGet:
def test_returns_none_when_no_conversation_id(self):
from application.api.answer.services.conversation_service import (
ConversationService,
)
assert ConversationService().get_conversation("", "u") is None
def test_returns_none_when_no_conversation_id(self, mock_mongo_db):
def test_returns_none_when_no_user_id(self):
from application.api.answer.services.conversation_service import (
ConversationService,
)
assert (
ConversationService().get_conversation(
"00000000-0000-0000-0000-000000000001", ""
)
is None
)
def test_returns_none_when_not_found(self, pg_conn):
from application.api.answer.services.conversation_service import (
ConversationService,
)
with _patch_db(pg_conn):
got = ConversationService().get_conversation(
"00000000-0000-0000-0000-000000000000", "u",
)
assert got is None
def test_returns_conversation_with_messages(self, pg_conn):
from application.api.answer.services.conversation_service import (
ConversationService,
)
from application.storage.db.repositories.conversations import (
ConversationsRepository,
)
user = "u-get"
repo = ConversationsRepository(pg_conn)
conv = repo.create(user, name="hi")
conv_id = str(conv["id"])
repo.append_message(
conv_id, {"prompt": "q1", "response": "r1"}
)
with _patch_db(pg_conn):
got = ConversationService().get_conversation(conv_id, user)
assert got is not None
assert got["_id"] == conv_id
assert len(got["queries"]) == 1
def test_handles_exception_returns_none(self):
from application.api.answer.services.conversation_service import (
ConversationService,
)
service = ConversationService()
result = service.get_conversation("", "user_123")
@contextmanager
def _broken():
raise RuntimeError("boom")
yield
assert result is None
def test_returns_none_when_no_user_id(self, mock_mongo_db):
from application.api.answer.services.conversation_service import (
ConversationService,
)
service = ConversationService()
result = service.get_conversation(str(ObjectId()), "")
assert result is None
def test_returns_conversation_for_owner(self, mock_mongo_db):
from application.api.answer.services.conversation_service import (
ConversationService,
)
from application.core.settings import settings
service = ConversationService()
collection = mock_mongo_db[settings.MONGO_DB_NAME]["conversations"]
conv_id = ObjectId()
conversation = {
"_id": conv_id,
"user": "user_123",
"name": "Test Conv",
"queries": [],
}
collection.insert_one(conversation)
result = service.get_conversation(str(conv_id), "user_123")
assert result is not None
assert result["name"] == "Test Conv"
assert result["_id"] == str(conv_id)
def test_returns_none_for_unauthorized_user(self, mock_mongo_db):
from application.api.answer.services.conversation_service import (
ConversationService,
)
from application.core.settings import settings
service = ConversationService()
collection = mock_mongo_db[settings.MONGO_DB_NAME]["conversations"]
conv_id = ObjectId()
collection.insert_one(
{"_id": conv_id, "user": "owner_123", "name": "Private Conv"}
)
result = service.get_conversation(str(conv_id), "hacker_456")
assert result is None
def test_converts_objectid_to_string(self, mock_mongo_db):
from application.api.answer.services.conversation_service import (
ConversationService,
)
from application.core.settings import settings
service = ConversationService()
collection = mock_mongo_db[settings.MONGO_DB_NAME]["conversations"]
conv_id = ObjectId()
collection.insert_one({"_id": conv_id, "user": "user_123", "name": "Test"})
result = service.get_conversation(str(conv_id), "user_123")
assert isinstance(result["_id"], str)
assert result["_id"] == str(conv_id)
with patch(
"application.api.answer.services.conversation_service.db_readonly",
_broken,
):
got = ConversationService().get_conversation("abc", "u")
assert got is None
@pytest.mark.unit
class TestConversationServiceSave:
def test_raises_error_when_no_user_in_token(self, mock_mongo_db):
"""Test validation: user ID required"""
def test_raises_for_none_token(self):
from application.api.answer.services.conversation_service import (
ConversationService,
)
service = ConversationService()
mock_llm = Mock()
with pytest.raises(ValueError, match="User ID not found"):
service.save_conversation(
with pytest.raises(ValueError):
ConversationService().save_conversation(
conversation_id=None,
question="Test?",
response="Answer",
thought="",
sources=[],
tool_calls=[],
llm=mock_llm,
model_id="gpt-4",
decoded_token={}, # No 'sub' key
question="q", response="r", thought="", sources=[],
tool_calls=[], llm=MagicMock(), model_id="gpt",
decoded_token=None,
)
def test_truncates_long_source_text(self, mock_mongo_db):
def test_raises_when_no_user_in_token(self):
from application.api.answer.services.conversation_service import (
ConversationService,
)
from application.core.settings import settings
from bson import ObjectId
service = ConversationService()
collection = mock_mongo_db[settings.MONGO_DB_NAME]["conversations"]
mock_llm = Mock()
mock_llm.gen.return_value = "Test Summary"
long_text = "x" * 2000
sources = [{"text": long_text, "title": "Doc"}]
conv_id = service.save_conversation(
conversation_id=None,
question="Question",
response="Response",
thought="",
sources=sources,
tool_calls=[],
llm=mock_llm,
model_id="gpt-4",
decoded_token={"sub": "user_123"},
)
saved_conv = collection.find_one({"_id": ObjectId(conv_id)})
saved_source_text = saved_conv["queries"][0]["sources"][0]["text"]
assert len(saved_source_text) == 1000
assert saved_source_text == "x" * 1000
def test_creates_new_conversation_with_summary(self, mock_mongo_db):
from application.api.answer.services.conversation_service import (
ConversationService,
)
from application.core.settings import settings
from bson import ObjectId
service = ConversationService()
collection = mock_mongo_db[settings.MONGO_DB_NAME]["conversations"]
mock_llm = Mock()
mock_llm.gen.return_value = "Python Basics"
conv_id = service.save_conversation(
conversation_id=None,
question="What is Python?",
response="Python is a programming language",
thought="",
sources=[],
tool_calls=[],
llm=mock_llm,
model_id="gpt-4",
decoded_token={"sub": "user_123"},
)
assert conv_id is not None
saved_conv = collection.find_one({"_id": ObjectId(conv_id)})
assert saved_conv["name"] == "Python Basics"
assert saved_conv["user"] == "user_123"
assert len(saved_conv["queries"]) == 1
assert saved_conv["queries"][0]["prompt"] == "What is Python?"
def test_appends_to_existing_conversation(self, mock_mongo_db):
from application.api.answer.services.conversation_service import (
ConversationService,
)
from application.core.settings import settings
from bson import ObjectId
service = ConversationService()
collection = mock_mongo_db[settings.MONGO_DB_NAME]["conversations"]
existing_conv_id = ObjectId()
collection.insert_one(
{
"_id": existing_conv_id,
"user": "user_123",
"name": "Old Conv",
"queries": [{"prompt": "Q1", "response": "A1"}],
}
)
mock_llm = Mock()
result = service.save_conversation(
conversation_id=str(existing_conv_id),
question="Q2",
response="A2",
thought="",
sources=[],
tool_calls=[],
llm=mock_llm,
model_id="gpt-4",
decoded_token={"sub": "user_123"},
)
assert result == str(existing_conv_id)
def test_prevents_unauthorized_conversation_update(self, mock_mongo_db):
from application.api.answer.services.conversation_service import (
ConversationService,
)
from application.core.settings import settings
service = ConversationService()
collection = mock_mongo_db[settings.MONGO_DB_NAME]["conversations"]
conv_id = ObjectId()
collection.insert_one({"_id": conv_id, "user": "owner_123", "queries": []})
mock_llm = Mock()
with pytest.raises(ValueError, match="not found or unauthorized"):
service.save_conversation(
conversation_id=str(conv_id),
question="Hack",
response="Attempt",
thought="",
sources=[],
tool_calls=[],
llm=mock_llm,
model_id="gpt-4",
decoded_token={"sub": "hacker_456"},
with pytest.raises(ValueError):
ConversationService().save_conversation(
conversation_id=None,
question="q", response="r", thought="", sources=[],
tool_calls=[], llm=MagicMock(), model_id="gpt",
decoded_token={},
)
def test_creates_new_conversation(self, pg_conn):
from application.api.answer.services.conversation_service import (
ConversationService,
)
from application.storage.db.repositories.conversations import (
ConversationsRepository,
)
user = "u-save-new"
mock_llm = MagicMock()
mock_llm.gen.return_value = "Title"
with _patch_db(pg_conn):
conv_id = ConversationService().save_conversation(
conversation_id=None,
question="q1", response="r1", thought="",
sources=[{"text": "x" * 2000}], tool_calls=[],
llm=mock_llm, model_id="gpt-4",
decoded_token={"sub": user},
)
got = ConversationsRepository(pg_conn).get_any(conv_id, user)
assert got is not None
assert got["name"] == "Title"
def test_appends_message_to_existing(self, pg_conn):
from application.api.answer.services.conversation_service import (
ConversationService,
)
from application.storage.db.repositories.conversations import (
ConversationsRepository,
)
user = "u-append"
repo = ConversationsRepository(pg_conn)
conv = repo.create(user, name="existing")
conv_id = str(conv["id"])
with _patch_db(pg_conn):
got = ConversationService().save_conversation(
conversation_id=conv_id,
question="q-new", response="r-new", thought="",
sources=[], tool_calls=[],
llm=MagicMock(), model_id="gpt-4",
decoded_token={"sub": user},
)
assert got == conv_id
messages = repo.get_messages(conv_id)
assert any(m["prompt"] == "q-new" for m in messages)
def test_updates_message_at_index(self, pg_conn):
from application.api.answer.services.conversation_service import (
ConversationService,
)
from application.storage.db.repositories.conversations import (
ConversationsRepository,
)
user = "u-update-at"
repo = ConversationsRepository(pg_conn)
conv = repo.create(user, name="edit")
conv_id = str(conv["id"])
repo.append_message(conv_id, {"prompt": "old", "response": "old"})
with _patch_db(pg_conn):
got = ConversationService().save_conversation(
conversation_id=conv_id,
question="updated", response="reply", thought="",
sources=[], tool_calls=[],
llm=MagicMock(), model_id="gpt-4",
decoded_token={"sub": user},
index=0,
)
assert got == conv_id
messages = repo.get_messages(conv_id)
assert messages[0]["prompt"] == "updated"
def test_raises_when_conversation_missing(self, pg_conn):
from application.api.answer.services.conversation_service import (
ConversationService,
)
with _patch_db(pg_conn), pytest.raises(ValueError):
ConversationService().save_conversation(
conversation_id="00000000-0000-0000-0000-000000000000",
question="q", response="r", thought="",
sources=[], tool_calls=[],
llm=MagicMock(), model_id="gpt",
decoded_token={"sub": "u"},
)
def test_save_with_empty_llm_title_falls_back(self, pg_conn):
from application.api.answer.services.conversation_service import (
ConversationService,
)
from application.storage.db.repositories.conversations import (
ConversationsRepository,
)
user = "u-empty-title"
mock_llm = MagicMock()
mock_llm.gen.return_value = ""
with _patch_db(pg_conn):
conv_id = ConversationService().save_conversation(
conversation_id=None,
question="q-fallback", response="r", thought="",
sources=[], tool_calls=[],
llm=mock_llm, model_id="gpt-4",
decoded_token={"sub": user},
)
got = ConversationsRepository(pg_conn).get_any(conv_id, user)
assert got["name"] == "q-fallback"
class TestCompressionMetadata:
def test_update_compression_metadata(self, pg_conn):
from application.api.answer.services.conversation_service import (
ConversationService,
)
from application.storage.db.repositories.conversations import (
ConversationsRepository,
)
user = "u-compress"
repo = ConversationsRepository(pg_conn)
conv = repo.create(user, name="c")
conv_id = str(conv["id"])
with _patch_db(pg_conn):
ConversationService().update_compression_metadata(
conv_id,
{
"timestamp": datetime.now(timezone.utc),
"compressed_summary": "summary",
"model_used": "gpt",
},
)
def test_update_compression_raises_on_error(self):
from application.api.answer.services.conversation_service import (
ConversationService,
)
@contextmanager
def _broken():
raise RuntimeError("boom")
yield
with patch(
"application.api.answer.services.conversation_service.db_session",
_broken,
), pytest.raises(RuntimeError):
ConversationService().update_compression_metadata("abc", {})
def test_append_compression_message_skips_empty_summary(self, pg_conn):
from application.api.answer.services.conversation_service import (
ConversationService,
)
with _patch_db(pg_conn):
ConversationService().append_compression_message(
"any-id", {"compressed_summary": ""}
)
def test_append_compression_message_appends_summary(self, pg_conn):
from application.api.answer.services.conversation_service import (
ConversationService,
)
from application.storage.db.repositories.conversations import (
ConversationsRepository,
)
user = "u-append-cs"
repo = ConversationsRepository(pg_conn)
conv = repo.create(user, name="c")
conv_id = str(conv["id"])
with _patch_db(pg_conn):
ConversationService().append_compression_message(
conv_id,
{
"compressed_summary": "A summary",
"timestamp": datetime.now(timezone.utc),
"model_used": "gpt-4",
},
)
messages = repo.get_messages(conv_id)
assert any(m["response"] == "A summary" for m in messages)
def test_append_compression_message_swallows_error(self):
from application.api.answer.services.conversation_service import (
ConversationService,
)
@contextmanager
def _broken():
raise RuntimeError("boom")
yield
with patch(
"application.api.answer.services.conversation_service.db_session",
_broken,
):
ConversationService().append_compression_message(
"abc", {"compressed_summary": "x"}
)
def test_get_compression_metadata_returns_none_for_missing(
self, pg_conn,
):
from application.api.answer.services.conversation_service import (
ConversationService,
)
with _patch_db(pg_conn):
got = ConversationService().get_compression_metadata(
"00000000-0000-0000-0000-000000000000"
)
assert got is None
def test_get_compression_metadata_non_uuid_does_not_raise(
self, pg_conn, caplog,
):
"""Non-UUID ids (legacy Mongo ObjectIds with no legacy_mongo_id row)
must return None without hitting ``CAST(:id AS uuid)`` — which
would raise and pollute logs with a stack trace every call."""
import logging
from application.api.answer.services.conversation_service import (
ConversationService,
)
with caplog.at_level(logging.ERROR):
with _patch_db(pg_conn):
got = ConversationService().get_compression_metadata(
"507f1f77bcf86cd799439011"
)
assert got is None
assert not any(
"Error getting compression metadata" in r.message for r in caplog.records
)
def test_get_compression_metadata_handles_exception(self):
from application.api.answer.services.conversation_service import (
ConversationService,
)
@contextmanager
def _broken():
raise RuntimeError("boom")
yield
with patch(
"application.api.answer.services.conversation_service.db_readonly",
_broken,
):
got = ConversationService().get_compression_metadata("abc")
assert got is None
tests/api/answer/services/test_prompt_renderer.py
+48 -48
View File
@@ -652,9 +652,13 @@ class TestPromptRendererIntegration:
@pytest.mark.unit
@pytest.mark.skip(
reason="Uses removed MongoDB.get_client() + user_tools mongo collection; "
"needs rewrite against UserToolsRepository / pg_conn."
)
class TestStreamProcessorPromptRendering:
def test_stream_processor_pre_fetch_docs_none_doc_mode(self, mock_mongo_db):
def test_stream_processor_pre_fetch_docs_none_doc_mode(self):
from application.api.answer.services.stream_processor import StreamProcessor
request_data = {"question": "Test question", "isNoneDoc": True}
@@ -665,7 +669,7 @@ class TestStreamProcessorPromptRendering:
assert docs_together is None
assert docs_list is None
def test_pre_fetch_tools_disabled_globally(self, mock_mongo_db, monkeypatch):
def test_pre_fetch_tools_disabled_globally(self, monkeypatch):
from application.api.answer.services.stream_processor import StreamProcessor
from application.core.settings import settings
@@ -678,7 +682,7 @@ class TestStreamProcessorPromptRendering:
assert result is None
def test_pre_fetch_tools_disabled_per_request(self, mock_mongo_db):
def test_pre_fetch_tools_disabled_per_request(self):
from application.api.answer.services.stream_processor import StreamProcessor
request_data = {"question": "test", "disable_tool_prefetch": True}
@@ -692,18 +696,17 @@ class TestStreamProcessorPromptRendering:
from unittest.mock import MagicMock, patch
from application.api.answer.services.stream_processor import StreamProcessor
from application.core.mongo_db import MongoDB
from bson import ObjectId
from application.core.settings import settings
db = MongoDB.get_client()[list(MongoDB.get_client().keys())[0]]
tool_doc = {
"_id": ObjectId(),
"name": "memory",
"user": "user1",
"status": True,
"config": {},
}
db["user_tools"].insert_one(tool_doc)
db = mock_mongo_db[settings.MONGO_DB_NAME]
db["user_tools"].insert_one(
{
"name": "memory",
"user": "user1",
"status": True,
"config": {},
}
)
request_data = {"question": "test"}
processor = StreamProcessor(request_data, {"sub": "user1"})
@@ -730,18 +733,17 @@ class TestStreamProcessorPromptRendering:
from unittest.mock import MagicMock, patch
from application.api.answer.services.stream_processor import StreamProcessor
from application.core.mongo_db import MongoDB
from bson import ObjectId
from application.core.settings import settings
db = MongoDB.get_client()[list(MongoDB.get_client().keys())[0]]
tool_doc = {
"_id": ObjectId(),
"name": "memory",
"user": "user1",
"status": True,
"config": {},
}
db["user_tools"].insert_one(tool_doc)
db = mock_mongo_db[settings.MONGO_DB_NAME]
db["user_tools"].insert_one(
{
"name": "memory",
"user": "user1",
"status": True,
"config": {},
}
)
request_data = {"question": "test"}
processor = StreamProcessor(request_data, {"sub": "user1"})
@@ -768,7 +770,7 @@ class TestStreamProcessorPromptRendering:
assert "memory" in result
assert "memory_ls" in result["memory"]
def test_pre_fetch_tools_no_tools_configured(self, mock_mongo_db):
def test_pre_fetch_tools_no_tools_configured(self):
from application.api.answer.services.stream_processor import StreamProcessor
request_data = {"question": "test"}
@@ -782,18 +784,17 @@ class TestStreamProcessorPromptRendering:
from unittest.mock import MagicMock, patch
from application.api.answer.services.stream_processor import StreamProcessor
from application.core.mongo_db import MongoDB
from bson import ObjectId
from application.core.settings import settings
db = MongoDB.get_client()[list(MongoDB.get_client().keys())[0]]
tool_doc = {
"_id": ObjectId(),
"name": "memory",
"user": "user1",
"status": True,
"config": {},
}
db["user_tools"].insert_one(tool_doc)
db = mock_mongo_db[settings.MONGO_DB_NAME]
db["user_tools"].insert_one(
{
"name": "memory",
"user": "user1",
"status": True,
"config": {},
}
)
request_data = {"question": "test"}
processor = StreamProcessor(request_data, {"sub": "user1"})
@@ -823,18 +824,17 @@ class TestStreamProcessorPromptRendering:
from unittest.mock import MagicMock, patch
from application.api.answer.services.stream_processor import StreamProcessor
from application.core.mongo_db import MongoDB
from bson import ObjectId
from application.core.settings import settings
db = MongoDB.get_client()[list(MongoDB.get_client().keys())[0]]
tool_doc = {
"_id": ObjectId(),
"name": "memory",
"user": "user1",
"status": True,
"config": {},
}
db["user_tools"].insert_one(tool_doc)
db = mock_mongo_db[settings.MONGO_DB_NAME]
db["user_tools"].insert_one(
{
"name": "memory",
"user": "user1",
"status": True,
"config": {},
}
)
request_data = {"question": "test"}
processor = StreamProcessor(request_data, {"sub": "user1"})
tests/api/answer/services/test_stream_processor.py
+33 -713
View File
@@ -1,56 +1,35 @@
"""Tests for application/api/answer/services/stream_processor.py.
The previous suite was tightly coupled to Mongo (mock_mongo_db fixture,
bson.ObjectId, bson.DBRef, find_one, etc.) which no longer exist after the
Mongo -> Postgres cutover. Rewriting these ~18 tests against the new
repositories (AgentsRepository / PromptsRepository / ConversationsRepository)
requires meaningful setup that is best done alongside the migration of the
StreamProcessor internals themselves.
"""
import pytest
from bson import DBRef, ObjectId
# A static 24-hex-char string that is a valid ObjectId hex representation.
_STATIC_OID = "507f1f77bcf86cd799439011"
@pytest.mark.unit
class TestGetPromptFunction:
def test_loads_custom_prompt_from_database(self, mock_mongo_db):
from application.api.answer.services.stream_processor import get_prompt
from application.core.settings import settings
prompts_collection = mock_mongo_db[settings.MONGO_DB_NAME]["prompts"]
prompt_id = ObjectId()
prompts_collection.insert_one(
{
"_id": prompt_id,
"content": "Custom prompt from database",
"user": "user_123",
}
)
result = get_prompt(str(prompt_id), prompts_collection)
assert result == "Custom prompt from database"
def test_raises_error_for_invalid_prompt_id(self, mock_mongo_db):
from application.api.answer.services.stream_processor import get_prompt
from application.core.settings import settings
prompts_collection = mock_mongo_db[settings.MONGO_DB_NAME]["prompts"]
with pytest.raises(ValueError, match="Invalid prompt ID"):
get_prompt(str(ObjectId()), prompts_collection)
def test_raises_error_for_malformed_id(self, mock_mongo_db):
from application.api.answer.services.stream_processor import get_prompt
from application.core.settings import settings
prompts_collection = mock_mongo_db[settings.MONGO_DB_NAME]["prompts"]
with pytest.raises(ValueError, match="Invalid prompt ID"):
get_prompt("not_a_valid_id", prompts_collection)
pass
@pytest.mark.unit
class TestStreamProcessorInitialization:
pass
def test_initializes_with_decoded_token(self, mock_mongo_db):
def test_initializes_with_decoded_token(self):
from application.api.answer.services.stream_processor import StreamProcessor
conv_id = _STATIC_OID
request_data = {
"question": "What is Python?",
"conversation_id": str(ObjectId()),
"conversation_id": conv_id,
}
decoded_token = {"sub": "user_123", "email": "test@example.com"}
@@ -61,7 +40,7 @@ class TestStreamProcessorInitialization:
assert processor.initial_user_id == "user_123"
assert processor.conversation_id == request_data["conversation_id"]
def test_initializes_without_token(self, mock_mongo_db):
def test_initializes_without_token(self):
from application.api.answer.services.stream_processor import StreamProcessor
request_data = {"question": "Test question"}
@@ -72,7 +51,7 @@ class TestStreamProcessorInitialization:
assert processor.initial_user_id is None
assert processor.data == request_data
def test_initializes_default_attributes(self, mock_mongo_db):
def test_initializes_default_attributes(self):
from application.api.answer.services.stream_processor import StreamProcessor
processor = StreamProcessor({"question": "Test"}, {"sub": "user_123"})
@@ -86,10 +65,10 @@ class TestStreamProcessorInitialization:
assert processor.is_shared_usage is False
assert processor.shared_token is None
def test_extracts_conversation_id_from_request(self, mock_mongo_db):
def test_extracts_conversation_id_from_request(self):
from application.api.answer.services.stream_processor import StreamProcessor
conv_id = str(ObjectId())
conv_id = _STATIC_OID
request_data = {"question": "Test", "conversation_id": conv_id}
processor = StreamProcessor(request_data, {"sub": "user_123"})
@@ -99,66 +78,9 @@ class TestStreamProcessorInitialization:
@pytest.mark.unit
class TestStreamProcessorHistoryLoading:
pass
def test_loads_history_from_existing_conversation(self, mock_mongo_db):
from application.api.answer.services.stream_processor import StreamProcessor
from application.core.settings import settings
conversations_collection = mock_mongo_db[settings.MONGO_DB_NAME][
"conversations"
]
conv_id = ObjectId()
conversations_collection.insert_one(
{
"_id": conv_id,
"user": "user_123",
"name": "Test Conv",
"queries": [
{"prompt": "What is Python?", "response": "Python is a language"},
{"prompt": "Tell me more", "response": "Python is versatile"},
],
}
)
request_data = {
"question": "How to install it?",
"conversation_id": str(conv_id),
}
processor = StreamProcessor(request_data, {"sub": "user_123"})
processor._load_conversation_history()
assert len(processor.history) == 2
assert processor.history[0]["prompt"] == "What is Python?"
assert processor.history[1]["response"] == "Python is versatile"
def test_raises_error_for_unauthorized_conversation(self, mock_mongo_db):
from application.api.answer.services.stream_processor import StreamProcessor
from application.core.settings import settings
conversations_collection = mock_mongo_db[settings.MONGO_DB_NAME][
"conversations"
]
conv_id = ObjectId()
conversations_collection.insert_one(
{
"_id": conv_id,
"user": "owner_123",
"name": "Private Conv",
"queries": [],
}
)
request_data = {"question": "Hack attempt", "conversation_id": str(conv_id)}
processor = StreamProcessor(request_data, {"sub": "hacker_456"})
with pytest.raises(ValueError, match="Conversation not found or unauthorized"):
processor._load_conversation_history()
def test_uses_request_history_when_no_conversation_id(self, mock_mongo_db):
def test_uses_request_history_when_no_conversation_id(self):
from application.api.answer.services.stream_processor import StreamProcessor
request_data = {
@@ -173,37 +95,9 @@ class TestStreamProcessorHistoryLoading:
@pytest.mark.unit
class TestStreamProcessorAgentConfiguration:
pass
def test_configures_agent_from_valid_api_key(self, mock_mongo_db):
from application.api.answer.services.stream_processor import StreamProcessor
from application.core.settings import settings
agents_collection = mock_mongo_db[settings.MONGO_DB_NAME]["agents"]
agent_id = ObjectId()
agents_collection.insert_one(
{
"_id": agent_id,
"key": "test_api_key_123",
"endpoint": "openai",
"model": "gpt-4",
"prompt_id": "default",
"user": "user_123",
}
)
request_data = {"question": "Test", "api_key": "test_api_key_123"}
processor = StreamProcessor(request_data, None)
try:
processor._configure_agent()
assert processor.agent_config is not None
assert processor.agent_id == str(agent_id)
except Exception as e:
assert "Invalid API Key" in str(e)
def test_uses_default_config_without_api_key(self, mock_mongo_db):
def test_uses_default_config_without_api_key(self):
from application.api.answer.services.stream_processor import StreamProcessor
request_data = {"question": "Test"}
@@ -214,152 +108,13 @@ class TestStreamProcessorAgentConfiguration:
assert isinstance(processor.agent_config, dict)
assert processor.agent_id is None
def test_conversation_agent_overrides_request_active_docs(self, mock_mongo_db):
from application.api.answer.services.stream_processor import StreamProcessor
from application.core.settings import settings
db = mock_mongo_db[settings.MONGO_DB_NAME]
agents_collection = db["agents"]
sources_collection = db["sources"]
conversations_collection = db["conversations"]
agent_id = ObjectId()
conversation_id = ObjectId()
agent_source_id = ObjectId()
request_source_id = ObjectId()
sources_collection.insert_many(
[
{"_id": agent_source_id, "name": "Agent source", "retriever": "classic"},
{
"_id": request_source_id,
"name": "Request source",
"retriever": "hybrid",
},
]
)
agents_collection.insert_one(
{
"_id": agent_id,
"key": "agent_key_2",
"user": "user_123",
"prompt_id": "default",
"agent_type": "classic",
"source": DBRef("sources", agent_source_id),
}
)
conversations_collection.insert_one(
{
"_id": conversation_id,
"user": "user_123",
"agent_id": str(agent_id),
"queries": [],
}
)
processor = StreamProcessor(
{
"question": "Test",
"conversation_id": str(conversation_id),
"active_docs": str(request_source_id),
},
{"sub": "user_123"},
)
processor._configure_agent()
processor._configure_source()
assert processor.agent_id == str(agent_id)
assert processor.source["active_docs"] == str(agent_source_id)
@pytest.mark.unit
class TestStreamProcessorDocPrefetch:
pass
def test_prefetch_not_skipped_for_agent_when_isNoneDoc_true(self, mock_mongo_db):
from unittest.mock import MagicMock
from application.api.answer.services.stream_processor import StreamProcessor
from application.core.settings import settings
db = mock_mongo_db[settings.MONGO_DB_NAME]
agents_collection = db["agents"]
sources_collection = db["sources"]
agent_id = ObjectId()
source_id = ObjectId()
sources_collection.insert_one(
{"_id": source_id, "name": "Agent source", "retriever": "classic"}
)
agents_collection.insert_one(
{
"_id": agent_id,
"key": "agent_prefetch_key",
"user": "user_123",
"prompt_id": "default",
"agent_type": "classic",
"source": DBRef("sources", source_id),
}
)
processor = StreamProcessor(
{
"question": "Summarize context",
"agent_id": str(agent_id),
"isNoneDoc": True,
},
{"sub": "user_123"},
)
processor.initialize()
mock_retriever = MagicMock()
mock_retriever.chunks = 2
mock_retriever.doc_token_limit = 50000
mock_retriever.search.return_value = [
{"text": "Agent doc content", "source": "agent.pdf"}
]
processor.create_retriever = MagicMock(return_value=mock_retriever)
docs_together, docs = processor.pre_fetch_docs("Summarize context")
processor.create_retriever.assert_called_once()
assert docs is not None
assert docs_together is not None
assert "Agent doc content" in docs_together
def test_configure_source_treats_default_string_as_no_docs(self, mock_mongo_db):
from application.api.answer.services.stream_processor import StreamProcessor
from application.core.settings import settings
db = mock_mongo_db[settings.MONGO_DB_NAME]
agents_collection = db["agents"]
agent_id = ObjectId()
agents_collection.insert_one(
{
"_id": agent_id,
"key": "agent_default_source_key",
"user": "user_123",
"prompt_id": "default",
"agent_type": "classic",
"source": "default",
}
)
processor = StreamProcessor(
{"question": "Hi", "api_key": "agent_default_source_key"},
None,
)
processor._configure_agent()
processor._configure_source()
assert processor.source == {}
assert processor.all_sources == []
def test_prefetch_skipped_when_no_active_docs(self, mock_mongo_db):
def test_prefetch_skipped_when_no_active_docs(self):
from unittest.mock import MagicMock
from application.api.answer.services.stream_processor import StreamProcessor
@@ -377,7 +132,7 @@ class TestStreamProcessorDocPrefetch:
assert docs_together is None
assert docs is None
def test_prefetch_skipped_when_active_docs_is_default(self, mock_mongo_db):
def test_prefetch_skipped_when_active_docs_is_default(self):
from unittest.mock import MagicMock
from application.api.answer.services.stream_processor import StreamProcessor
@@ -395,132 +150,15 @@ class TestStreamProcessorDocPrefetch:
assert docs_together is None
assert docs is None
def test_agent_retriever_and_chunks_propagate_to_retriever_config(self, mock_mongo_db):
from application.api.answer.services.stream_processor import StreamProcessor
from application.core.settings import settings
db = mock_mongo_db[settings.MONGO_DB_NAME]
agents_collection = db["agents"]
source_id = ObjectId()
db["sources"].insert_one(
{"_id": source_id, "name": "src", "retriever": "hybrid", "chunks": 5}
)
agent_id = ObjectId()
agents_collection.insert_one(
{
"_id": agent_id,
"key": "agent_ret_key",
"user": "user_123",
"prompt_id": "default",
"agent_type": "classic",
"retriever": "hybrid",
"chunks": 5,
"source": DBRef("sources", source_id),
}
)
processor = StreamProcessor(
{"question": "Test", "api_key": "agent_ret_key"},
None,
)
processor.initialize()
assert processor.retriever_config["retriever_name"] == "hybrid"
assert processor.retriever_config["chunks"] == 5
def test_request_retriever_and_chunks_override_agent_config(self, mock_mongo_db):
from application.api.answer.services.stream_processor import StreamProcessor
from application.core.settings import settings
db = mock_mongo_db[settings.MONGO_DB_NAME]
agents_collection = db["agents"]
agent_id = ObjectId()
agents_collection.insert_one(
{
"_id": agent_id,
"key": "agent_override_key",
"user": "user_123",
"prompt_id": "default",
"agent_type": "classic",
"retriever": "hybrid",
"chunks": 5,
}
)
processor = StreamProcessor(
{
"question": "Test",
"api_key": "agent_override_key",
"retriever": "classic",
"chunks": 7,
},
None,
)
processor.initialize()
assert processor.retriever_config["retriever_name"] == "classic"
assert processor.retriever_config["chunks"] == 7
def test_agent_data_fetched_once_per_request(self, mock_mongo_db):
from unittest.mock import patch
from application.api.answer.services.stream_processor import StreamProcessor
from application.core.settings import settings
db = mock_mongo_db[settings.MONGO_DB_NAME]
agents_collection = db["agents"]
agent_id = ObjectId()
agents_collection.insert_one(
{
"_id": agent_id,
"key": "agent_once_key",
"user": "user_123",
"prompt_id": "default",
"agent_type": "classic",
}
)
processor = StreamProcessor(
{"question": "Test", "api_key": "agent_once_key"},
None,
)
with patch.object(
processor, "_get_data_from_api_key", wraps=processor._get_data_from_api_key
) as spy:
processor.initialize()
assert spy.call_count == 1
@pytest.mark.unit
class TestStreamProcessorAttachments:
pass
def test_processes_attachments_from_request(self, mock_mongo_db):
from application.api.answer.services.stream_processor import StreamProcessor
from application.core.settings import settings
attachments_collection = mock_mongo_db[settings.MONGO_DB_NAME]["attachments"]
att_id = ObjectId()
attachments_collection.insert_one(
{
"_id": att_id,
"filename": "document.pdf",
"content": "Document content",
"user": "user_123",
}
)
request_data = {"question": "Analyze this", "attachments": [str(att_id)]}
processor = StreamProcessor(request_data, {"sub": "user_123"})
assert processor.data.get("attachments") == [str(att_id)]
def test_handles_empty_attachments(self, mock_mongo_db):
def test_handles_empty_attachments(self):
from application.api.answer.services.stream_processor import StreamProcessor
request_data = {"question": "Simple question"}
@@ -536,326 +174,8 @@ class TestStreamProcessorAttachments:
@pytest.mark.unit
class TestToolPreFetch:
"""Tests for tool pre-fetching with saved parameter values from MongoDB"""
"""Tests for tool pre-fetching with saved parameter values."""
def test_cryptoprice_prefetch_with_saved_parameters(self, mock_mongo_db):
"""Test that cryptoprice tool is pre-fetched with saved parameter values from MongoDB structure"""
from application.api.answer.services.stream_processor import StreamProcessor
from application.core.settings import settings
from unittest.mock import patch, MagicMock
# Setup MongoDB with cryptoprice tool configuration
# NOTE: The collection is called "user_tools" not "tools"
tools_collection = mock_mongo_db[settings.MONGO_DB_NAME]["user_tools"]
tool_id = ObjectId()
tools_collection.insert_one(
{
"_id": tool_id,
"name": "cryptoprice",
"user": "user_123",
"status": True, # Must be True for tool to be included
"actions": [
{
"name": "cryptoprice_get",
"description": "Get cryptocurrency price",
"parameters": {
"type": "object",
"properties": {
"symbol": {
"type": "string",
"description": "Crypto symbol",
"value": "BTC" # Saved value in MongoDB
},
"currency": {
"type": "string",
"description": "Currency for price",
"value": "USD" # Saved value in MongoDB
}
},
"required": ["symbol", "currency"]
}
}
],
"config": {
"token": ""
}
}
)
request_data = {
"question": "What is the price of Bitcoin?",
"tools": [str(tool_id)]
}
processor = StreamProcessor(request_data, {"sub": "user_123"})
processor._required_tool_actions = {"cryptoprice": {"cryptoprice_get"}}
# Mock the ToolManager and tool instance
with patch("application.agents.tools.tool_manager.ToolManager") as mock_manager_class:
mock_manager = MagicMock()
mock_manager_class.return_value = mock_manager
# Mock the tool instance returned by load_tool
mock_tool = MagicMock()
mock_manager.load_tool.return_value = mock_tool
# Mock get_actions_metadata on the tool instance
mock_tool.get_actions_metadata.return_value = [
{
"name": "cryptoprice_get",
"description": "Get cryptocurrency price",
"parameters": {
"type": "object",
"properties": {
"symbol": {"type": "string", "description": "Crypto symbol"},
"currency": {"type": "string", "description": "Currency for price"}
},
"required": ["symbol", "currency"]
}
}
]
# Mock execute_action on the tool instance to return price data
mock_tool.execute_action.return_value = {
"status_code": 200,
"price": 45000.50,
"message": "Price of BTC in USD retrieved successfully."
}
# Execute pre-fetch
tools_data = processor.pre_fetch_tools()
# Verify the tool was called
assert mock_tool.execute_action.called
# Verify it was called with the saved parameters from MongoDB
call_args = mock_tool.execute_action.call_args
assert call_args is not None
# Check action name uses the full metadata name for execution
assert call_args[0][0] == "cryptoprice_get"
# Check kwargs contain saved values
kwargs = call_args[1]
assert kwargs.get("symbol") == "BTC"
assert kwargs.get("currency") == "USD"
# Verify tools_data structure
assert "cryptoprice" in tools_data
# Results are exposed under the full action name
assert "cryptoprice_get" in tools_data["cryptoprice"]
assert tools_data["cryptoprice"]["cryptoprice_get"]["price"] == 45000.50
def test_prefetch_with_missing_saved_values_uses_defaults(self, mock_mongo_db):
"""Test that pre-fetch falls back to defaults when saved values are missing"""
from application.api.answer.services.stream_processor import StreamProcessor
from application.core.settings import settings
from unittest.mock import patch, MagicMock
tools_collection = mock_mongo_db[settings.MONGO_DB_NAME]["user_tools"]
tool_id = ObjectId()
# Tool configuration without saved values
tools_collection.insert_one(
{
"_id": tool_id,
"name": "cryptoprice",
"user": "user_123",
"status": True,
"actions": [
{
"name": "cryptoprice_get",
"description": "Get cryptocurrency price",
"parameters": {
"type": "object",
"properties": {
"symbol": {
"type": "string",
"description": "Crypto symbol",
"default": "ETH" # Only default, no saved value
},
"currency": {
"type": "string",
"description": "Currency",
"default": "EUR"
}
},
"required": ["symbol", "currency"]
}
}
],
"config": {}
}
)
request_data = {
"question": "Crypto price?",
"tools": [str(tool_id)]
}
processor = StreamProcessor(request_data, {"sub": "user_123"})
processor._required_tool_actions = {"cryptoprice": {"cryptoprice_get"}}
with patch("application.agents.tools.tool_manager.ToolManager") as mock_manager_class:
mock_manager = MagicMock()
mock_manager_class.return_value = mock_manager
# Mock the tool instance
mock_tool = MagicMock()
mock_manager.load_tool.return_value = mock_tool
mock_tool.get_actions_metadata.return_value = [
{
"name": "cryptoprice_get",
"parameters": {
"type": "object",
"properties": {
"symbol": {"type": "string", "default": "ETH"},
"currency": {"type": "string", "default": "EUR"}
}
}
}
]
mock_tool.execute_action.return_value = {
"status_code": 200,
"price": 2500.00
}
processor.pre_fetch_tools()
# Should use default values when saved values are missing
call_args = mock_tool.execute_action.call_args
if call_args:
kwargs = call_args[1]
# Either uses defaults or skips if no values available
assert kwargs.get("symbol") in ["ETH", None]
assert kwargs.get("currency") in ["EUR", None]
def test_prefetch_with_tool_id_reference(self, mock_mongo_db):
"""Test that tools can be referenced by MongoDB ObjectId in templates"""
from application.api.answer.services.stream_processor import StreamProcessor
from application.core.settings import settings
from unittest.mock import patch, MagicMock
tools_collection = mock_mongo_db[settings.MONGO_DB_NAME]["user_tools"]
tool_id = ObjectId()
# Create a tool in the database
tools_collection.insert_one(
{
"_id": tool_id,
"name": "memory",
"user": "user_123",
"status": True,
"actions": [
{
"name": "memory_ls",
"description": "List files",
"parameters": {
"type": "object",
"properties": {}
}
}
],
"config": {},
}
)
request_data = {"question": "test"}
processor = StreamProcessor(request_data, {"sub": "user_123"})
# Mock the filtering to require this specific tool by ID
processor._required_tool_actions = {
str(tool_id): {"memory_ls"} # Reference by ObjectId string
}
with patch(
"application.agents.tools.tool_manager.ToolManager"
) as mock_manager_class:
mock_manager = MagicMock()
mock_manager_class.return_value = mock_manager
# Mock the tool instance
mock_tool = MagicMock()
mock_manager.load_tool.return_value = mock_tool
mock_tool.get_actions_metadata.return_value = [
{"name": "memory_ls", "description": "List files", "parameters": {"properties": {}}}
]
mock_tool.execute_action.return_value = "Directory: /\n- file.txt"
result = processor.pre_fetch_tools()
# Tool data should be available under both name and ID
assert result is not None
assert "memory" in result
assert str(tool_id) in result
# Both should point to the same data
assert result["memory"] == result[str(tool_id)]
assert "memory_ls" in result[str(tool_id)]
def test_prefetch_with_multiple_same_name_tools(self, mock_mongo_db):
"""Test that multiple tools with the same name can be distinguished by ID"""
from application.api.answer.services.stream_processor import StreamProcessor
from application.core.settings import settings
from unittest.mock import patch, MagicMock
tools_collection = mock_mongo_db[settings.MONGO_DB_NAME]["user_tools"]
# Create two memory tools with different IDs
tool_id_1 = ObjectId()
tool_id_2 = ObjectId()
tools_collection.insert_many([
{
"_id": tool_id_1,
"name": "memory",
"user": "user_123",
"status": True,
"actions": [{"name": "memory_ls", "parameters": {"properties": {}}}],
"config": {"path": "/home"},
},
{
"_id": tool_id_2,
"name": "memory",
"user": "user_123",
"status": True,
"actions": [{"name": "memory_ls", "parameters": {"properties": {}}}],
"config": {"path": "/work"},
}
])
request_data = {"question": "test"}
processor = StreamProcessor(request_data, {"sub": "user_123"})
# Mock the filtering to require only the second tool by ID
processor._required_tool_actions = {
str(tool_id_2): {"memory_ls"} # Only reference the second one
}
with patch(
"application.agents.tools.tool_manager.ToolManager"
) as mock_manager_class:
mock_manager = MagicMock()
mock_manager_class.return_value = mock_manager
# Mock the tool instance
mock_tool = MagicMock()
mock_manager.load_tool.return_value = mock_tool
mock_tool.get_actions_metadata.return_value = [
{"name": "memory_ls", "parameters": {"properties": {}}}
]
mock_tool.execute_action.return_value = "Work directory"
result = processor.pre_fetch_tools()
# Only the second tool should be fetched (referenced by ID)
assert result is not None
assert str(tool_id_2) in result
# Since filtering is enabled and only tool_id_2 is referenced,
# only tool_id_2 should be pre-fetched
# The "memory" key will still exist because we store under both name and ID
assert "memory" in result
tests/api/answer/services/test_stream_processor_utils.py
+582
View File
@@ -0,0 +1,582 @@
"""Targeted tests for application/api/answer/services/stream_processor.py.
Tests the ``get_prompt`` helper and simpler StreamProcessor methods against
real ephemeral Postgres.
"""
from contextlib import contextmanager
from unittest.mock import patch
import pytest
@contextmanager
def _patch_db(conn):
@contextmanager
def _yield():
yield conn
with patch(
"application.api.answer.services.stream_processor.db_readonly", _yield
), patch(
"application.api.answer.services.stream_processor.db_session", _yield
):
yield
class TestGetPrompt:
def test_default_preset(self):
from application.api.answer.services.stream_processor import get_prompt
got = get_prompt("default")
assert isinstance(got, str)
assert len(got) > 0
def test_creative_preset(self):
from application.api.answer.services.stream_processor import get_prompt
got = get_prompt("creative")
assert isinstance(got, str) and len(got) > 0
def test_strict_preset(self):
from application.api.answer.services.stream_processor import get_prompt
got = get_prompt("strict")
assert isinstance(got, str) and len(got) > 0
def test_agentic_default_preset(self):
from application.api.answer.services.stream_processor import get_prompt
got = get_prompt("agentic_default")
assert isinstance(got, str) and len(got) > 0
def test_none_defaults_to_default(self):
from application.api.answer.services.stream_processor import get_prompt
got = get_prompt(None)
assert isinstance(got, str) and len(got) > 0
def test_empty_string_defaults_to_default(self):
from application.api.answer.services.stream_processor import get_prompt
assert get_prompt("") == get_prompt("default")
def test_non_string_id_converted(self):
from application.api.answer.services.stream_processor import get_prompt
# A UUID object would be stringified; use an int to test the branch
with pytest.raises(ValueError):
# Int converts to str "42" which isn't a preset, and will
# raise ValueError once lookup fails
get_prompt(42)
def test_unknown_prompt_id_raises(self, pg_conn):
from application.api.answer.services.stream_processor import get_prompt
with _patch_db(pg_conn), pytest.raises(ValueError):
get_prompt("00000000-0000-0000-0000-000000000000")
def test_legacy_id_unknown_raises(self, pg_conn):
from application.api.answer.services.stream_processor import get_prompt
with _patch_db(pg_conn), pytest.raises(ValueError):
get_prompt("507f1f77bcf86cd799439011")
def test_uuid_lookup_returns_content(self, pg_conn):
from application.api.answer.services.stream_processor import get_prompt
from application.storage.db.repositories.prompts import (
PromptsRepository,
)
prompt = PromptsRepository(pg_conn).create(
"u", "myprompt", "custom prompt content",
)
with _patch_db(pg_conn):
got = get_prompt(str(prompt["id"]))
assert got == "custom prompt content"
class TestStreamProcessorInit:
def test_basic_init(self):
from application.api.answer.services.stream_processor import (
StreamProcessor,
)
data = {"question": "hi", "conversation_id": "conv-1"}
token = {"sub": "user-abc"}
sp = StreamProcessor(data, token)
assert sp.data == data
assert sp.decoded_token == token
assert sp.initial_user_id == "user-abc"
assert sp.conversation_id == "conv-1"
def test_init_no_token(self):
from application.api.answer.services.stream_processor import (
StreamProcessor,
)
sp = StreamProcessor({"question": "hi"}, None)
assert sp.decoded_token is None
assert sp.initial_user_id is None
assert sp.conversation_id is None
def test_init_sets_agent_id_from_data(self):
from application.api.answer.services.stream_processor import (
StreamProcessor,
)
data = {"question": "hi", "agent_id": "agent-xyz"}
sp = StreamProcessor(data, {"sub": "u"})
assert sp.agent_id == "agent-xyz"
class TestLoadConversationHistory:
def test_no_conversation_id_uses_request_history(self):
from application.api.answer.services.stream_processor import (
StreamProcessor,
)
import json as _json
data = {
"question": "hi",
"history": _json.dumps([{"prompt": "a", "response": "b"}]),
}
sp = StreamProcessor(data, {"sub": "u"})
sp._load_conversation_history()
assert len(sp.history) == 1
def test_loads_existing_conversation_history(self, pg_conn):
from application.api.answer.services.stream_processor import (
StreamProcessor,
)
from application.storage.db.repositories.conversations import (
ConversationsRepository,
)
user = "u-load-hist"
repo = ConversationsRepository(pg_conn)
conv = repo.create(user, name="c")
conv_id = str(conv["id"])
repo.append_message(
conv_id,
{"prompt": "q1", "response": "r1"},
)
repo.append_message(
conv_id,
{"prompt": "q2", "response": "r2"},
)
sp = StreamProcessor(
{"question": "x", "conversation_id": conv_id},
{"sub": user},
)
# Also patch conversation_service.get_conversation's DB accessor
with _patch_db(pg_conn), patch(
"application.api.answer.services.conversation_service.db_readonly",
) as mock_readonly:
@contextmanager
def _yield():
yield pg_conn
mock_readonly.side_effect = _yield
sp._load_conversation_history()
assert len(sp.history) == 2
def test_unauthorized_conversation_raises(self, pg_conn):
from application.api.answer.services.stream_processor import (
StreamProcessor,
)
from application.storage.db.repositories.conversations import (
ConversationsRepository,
)
repo = ConversationsRepository(pg_conn)
conv = repo.create("owner-user", name="c")
sp = StreamProcessor(
{"question": "hack", "conversation_id": str(conv["id"])},
{"sub": "hacker"},
)
with _patch_db(pg_conn), patch(
"application.api.answer.services.conversation_service.db_readonly"
) as mock_readonly:
@contextmanager
def _yield():
yield pg_conn
mock_readonly.side_effect = _yield
with pytest.raises(ValueError):
sp._load_conversation_history()
class TestHasActiveDocs:
def test_false_when_no_active_docs(self):
from application.api.answer.services.stream_processor import (
StreamProcessor,
)
sp = StreamProcessor({"question": "q"}, {"sub": "u"})
sp.source = {}
sp.all_sources = []
assert sp._has_active_docs() is False
def test_true_when_source_active(self):
from application.api.answer.services.stream_processor import (
StreamProcessor,
)
sp = StreamProcessor({"question": "q"}, {"sub": "u"})
sp.source = {"active_docs": "abc"}
assert sp._has_active_docs() is True
def test_false_when_source_empty(self):
from application.api.answer.services.stream_processor import (
StreamProcessor,
)
sp = StreamProcessor({"question": "q"}, {"sub": "u"})
sp.source = None
assert sp._has_active_docs() is False
def test_default_returns_false(self):
from application.api.answer.services.stream_processor import (
StreamProcessor,
)
sp = StreamProcessor({"question": "q"}, {"sub": "u"})
sp.source = {"active_docs": "default"}
sp.all_sources = []
assert sp._has_active_docs() is False
class TestProcessAttachments:
def test_no_attachments(self, pg_conn):
from application.api.answer.services.stream_processor import (
StreamProcessor,
)
sp = StreamProcessor({"question": "q"}, {"sub": "u"})
# attachment_ids not set
with _patch_db(pg_conn):
sp._process_attachments()
assert sp.attachments == []
def test_retrieves_attachments_by_id(self, pg_conn):
from application.api.answer.services.stream_processor import (
StreamProcessor,
)
from application.storage.db.repositories.attachments import (
AttachmentsRepository,
)
user = "u-atts"
att = AttachmentsRepository(pg_conn).create(
user, "doc.txt", "/path",
content="content here",
size=100,
)
sp = StreamProcessor(
{"question": "q", "attachments": [str(att["id"])]},
{"sub": user},
)
with _patch_db(pg_conn):
sp._process_attachments()
assert len(sp.attachments) == 1
assert sp.attachments[0]["content"] == "content here"
class TestGetAttachmentsContent:
def test_empty_list(self, pg_conn):
from application.api.answer.services.stream_processor import (
StreamProcessor,
)
sp = StreamProcessor({"question": "q"}, {"sub": "u"})
with _patch_db(pg_conn):
got = sp._get_attachments_content([], "u")
assert got == []
def test_skips_missing_attachments(self, pg_conn):
from application.api.answer.services.stream_processor import (
StreamProcessor,
)
sp = StreamProcessor({"question": "q"}, {"sub": "u"})
with _patch_db(pg_conn):
got = sp._get_attachments_content(
["00000000-0000-0000-0000-000000000000"], "u",
)
assert got == []
class TestResolveAgentId:
def test_returns_agent_id_from_request(self):
from application.api.answer.services.stream_processor import (
StreamProcessor,
)
sp = StreamProcessor({"agent_id": "req-agent"}, {"sub": "u"})
assert sp._resolve_agent_id() == "req-agent"
def test_returns_none_if_not_set(self, pg_conn):
from application.api.answer.services.stream_processor import (
StreamProcessor,
)
sp = StreamProcessor({}, {"sub": "u"})
assert sp._resolve_agent_id() is None
class TestGetAgentKey:
def test_returns_tuple_for_none_agent_id(self):
from application.api.answer.services.stream_processor import (
StreamProcessor,
)
sp = StreamProcessor({}, {"sub": "u"})
key, is_shared, tok = sp._get_agent_key(None, "u")
assert key is None and is_shared is False and tok is None
def test_raises_for_missing_agent(self, pg_conn):
from application.api.answer.services.stream_processor import (
StreamProcessor,
)
sp = StreamProcessor({}, {"sub": "u"})
with _patch_db(pg_conn), pytest.raises(Exception):
sp._get_agent_key(
"00000000-0000-0000-0000-000000000000", "u",
)
def test_returns_key_for_owned_agent(self, pg_conn):
from application.api.answer.services.stream_processor import (
StreamProcessor,
)
from application.storage.db.repositories.agents import AgentsRepository
agent = AgentsRepository(pg_conn).create(
"owner", "a", "published", key="the-key",
)
sp = StreamProcessor({}, {"sub": "owner"})
with _patch_db(pg_conn):
key, shared, tok = sp._get_agent_key(str(agent["id"]), "owner")
assert key == "the-key"
assert shared is False
def test_raises_on_unauthorized_access(self, pg_conn):
from application.api.answer.services.stream_processor import (
StreamProcessor,
)
from application.storage.db.repositories.agents import AgentsRepository
agent = AgentsRepository(pg_conn).create(
"owner", "a", "published", key="k", shared=False,
)
sp = StreamProcessor({}, {"sub": "not-owner"})
with _patch_db(pg_conn), pytest.raises(Exception):
sp._get_agent_key(str(agent["id"]), "not-owner")
class TestConfigureSource:
def test_agent_data_with_sources_list(self):
from application.api.answer.services.stream_processor import (
StreamProcessor,
)
sp = StreamProcessor({}, {"sub": "u"})
sp._agent_data = {
"sources": [
{"id": "s1", "retriever": "classic"},
{"id": "default"},
],
}
sp._configure_source()
assert sp.source == {"active_docs": ["s1"]}
def test_agent_data_with_single_source(self):
from application.api.answer.services.stream_processor import (
StreamProcessor,
)
sp = StreamProcessor({}, {"sub": "u"})
sp._agent_data = {"source": "src-1", "retriever": "classic"}
sp._configure_source()
assert sp.source == {"active_docs": "src-1"}
def test_agent_data_default_source(self):
from application.api.answer.services.stream_processor import (
StreamProcessor,
)
sp = StreamProcessor({}, {"sub": "u"})
sp._agent_data = {"source": "default"}
sp._configure_source()
assert sp.source == {}
def test_request_active_docs_used(self):
from application.api.answer.services.stream_processor import (
StreamProcessor,
)
sp = StreamProcessor({"active_docs": "abc"}, {"sub": "u"})
sp._configure_source()
assert sp.source == {"active_docs": "abc"}
def test_request_active_docs_default(self):
from application.api.answer.services.stream_processor import (
StreamProcessor,
)
sp = StreamProcessor({"active_docs": "default"}, {"sub": "u"})
sp._configure_source()
assert sp.source == {}
def test_no_data_empty_source(self):
from application.api.answer.services.stream_processor import (
StreamProcessor,
)
sp = StreamProcessor({}, {"sub": "u"})
sp._configure_source()
assert sp.source == {}
class TestConfigureRetriever:
def test_defaults(self):
from application.api.answer.services.stream_processor import (
StreamProcessor,
)
sp = StreamProcessor({}, {"sub": "u"})
sp._configure_retriever()
assert sp.retriever_config["retriever_name"] == "classic"
assert sp.retriever_config["chunks"] == 2
def test_agent_overrides(self):
from application.api.answer.services.stream_processor import (
StreamProcessor,
)
sp = StreamProcessor({}, {"sub": "u"})
sp._agent_data = {"retriever": "hybrid_search", "chunks": 5}
sp._configure_retriever()
assert sp.retriever_config["retriever_name"] == "hybrid_search"
assert sp.retriever_config["chunks"] == 5
def test_request_overrides_agent(self):
from application.api.answer.services.stream_processor import (
StreamProcessor,
)
sp = StreamProcessor(
{"retriever": "duckdb", "chunks": 7}, {"sub": "u"},
)
sp._agent_data = {"retriever": "hybrid_search", "chunks": 5}
sp._configure_retriever()
assert sp.retriever_config["retriever_name"] == "duckdb"
assert sp.retriever_config["chunks"] == 7
def test_invalid_agent_chunks_falls_back(self):
from application.api.answer.services.stream_processor import (
StreamProcessor,
)
sp = StreamProcessor({}, {"sub": "u"})
sp._agent_data = {"chunks": "not-a-number"}
sp._configure_retriever()
assert sp.retriever_config["chunks"] == 2
def test_invalid_request_chunks_falls_back(self):
from application.api.answer.services.stream_processor import (
StreamProcessor,
)
sp = StreamProcessor({"chunks": "abc"}, {"sub": "u"})
sp._configure_retriever()
assert sp.retriever_config["chunks"] == 2
def test_isnonedoc_without_api_key_sets_chunks_to_0(self):
from application.api.answer.services.stream_processor import (
StreamProcessor,
)
sp = StreamProcessor({"isNoneDoc": True}, {"sub": "u"})
sp.agent_key = None
sp._configure_retriever()
assert sp.retriever_config["chunks"] == 0
class TestGetPromptContent:
def test_gets_from_agent_config(self, pg_conn):
from application.api.answer.services.stream_processor import (
StreamProcessor,
)
from application.storage.db.repositories.prompts import (
PromptsRepository,
)
prompt = PromptsRepository(pg_conn).create(
"u", "p1", "My prompt content",
)
sp = StreamProcessor({}, {"sub": "u"})
sp.agent_config = {"prompt_id": str(prompt["id"])}
with _patch_db(pg_conn):
content = sp._get_prompt_content()
assert content == "My prompt content"
def test_returns_none_on_missing(self, pg_conn):
from application.api.answer.services.stream_processor import (
StreamProcessor,
)
sp = StreamProcessor({}, {"sub": "u"})
sp.agent_config = {
"prompt_id": "00000000-0000-0000-0000-000000000000",
}
with _patch_db(pg_conn):
content = sp._get_prompt_content()
assert content is None
def test_caches_prompt_content(self, pg_conn):
from application.api.answer.services.stream_processor import (
StreamProcessor,
)
sp = StreamProcessor({}, {"sub": "u"})
sp._prompt_content = "cached"
# Even with no agent_config, cached value returned
assert sp._get_prompt_content() == "cached"
class TestPreFetchDocs:
def test_skips_when_no_active_docs(self):
from application.api.answer.services.stream_processor import (
StreamProcessor,
)
sp = StreamProcessor({}, {"sub": "u"})
sp.source = {}
docs, raw = sp.pre_fetch_docs("question")
assert docs is None and raw is None
def test_skips_when_isnonedoc_no_agent(self):
from application.api.answer.services.stream_processor import (
StreamProcessor,
)
sp = StreamProcessor({"isNoneDoc": True}, {"sub": "u"})
sp.source = {"active_docs": "abc"} # would normally be active
sp.agent_id = None
docs, raw = sp.pre_fetch_docs("q")
assert docs is None and raw is None
def test_handles_retriever_exception(self):
from application.api.answer.services.stream_processor import (
StreamProcessor,
)
sp = StreamProcessor({}, {"sub": "u"})
sp.source = {"active_docs": "src"}
with patch(
"application.api.answer.services.stream_processor.StreamProcessor.create_retriever",
side_effect=RuntimeError("boom"),
):
docs, raw = sp.pre_fetch_docs("q")
assert docs is None and raw is None
class TestPreFetchTools:
def test_disabled_globally_returns_none(self):
from application.api.answer.services.stream_processor import (
StreamProcessor,
)
sp = StreamProcessor({}, {"sub": "u"})
with patch(
"application.api.answer.services.stream_processor.settings.ENABLE_TOOL_PREFETCH",
False,
):
got = sp.pre_fetch_tools()
assert got is None
def test_disabled_per_request(self):
from application.api.answer.services.stream_processor import (
StreamProcessor,
)
sp = StreamProcessor({"disable_tool_prefetch": True}, {"sub": "u"})
with patch(
"application.api.answer.services.stream_processor.settings.ENABLE_TOOL_PREFETCH",
True,
):
got = sp.pre_fetch_tools()
assert got is None
def test_no_user_tools_returns_none(self, pg_conn):
from application.api.answer.services.stream_processor import (
StreamProcessor,
)
sp = StreamProcessor({}, {"sub": "no-tools-user"})
with _patch_db(pg_conn), patch(
"application.api.answer.services.stream_processor.settings.ENABLE_TOOL_PREFETCH",
True,
):
got = sp.pre_fetch_tools()
assert got is None