| 1 | from __future__ import annotations |
| 2 | |
| 3 | import asyncio |
| 4 | import sys |
| 5 | from pathlib import Path |
| 6 | |
| 7 | PROJECT_ROOT = Path(__file__).resolve().parents[1] |
| 8 | if str(PROJECT_ROOT) not in sys.path: |
| 9 | sys.path.insert(0, str(PROJECT_ROOT)) |
| 10 | |
| 11 | from helpers.document_query import DocumentQueryHelper |
| 12 | |
| 13 | |
| 14 | class FakeStore: |
| 15 | @staticmethod |
| 16 | def normalize_uri(uri: str) -> str: |
| 17 | return uri |
| 18 | |
| 19 | async def search_documents(self, **_kwargs): |
| 20 | return [] |
| 21 | |
| 22 | |
| 23 | class FakeAgent: |
| 24 | def __init__(self): |
| 25 | self.chat_messages = None |
| 26 | |
| 27 | async def handle_intervention(self): |
| 28 | return None |
| 29 | |
| 30 | def parse_prompt(self, name: str) -> str: |
| 31 | return name |
| 32 | |
| 33 | async def call_utility_model(self, **_kwargs) -> str: |
| 34 | return "codename" |
| 35 | |
| 36 | async def call_chat_model(self, messages, explicit_caching=False): |
| 37 | self.chat_messages = messages |
| 38 | return "The project codename is Atlas.", None |
| 39 | |
| 40 | |
| 41 | def test_document_qa_uses_small_document_content_when_search_finds_no_chunks(): |
| 42 | agent = FakeAgent() |
| 43 | progress = [] |
| 44 | helper = object.__new__(DocumentQueryHelper) |
| 45 | helper.agent = agent |
| 46 | helper.store = FakeStore() |
| 47 | helper.config = {} |
| 48 | helper.progress_callback = progress.append |
| 49 | |
| 50 | async def document_get_content(uri, add_to_db=False): |
| 51 | assert uri == "/tmp/project.md" |
| 52 | assert add_to_db is True |
| 53 | return "# Project\n\nCodename: Atlas\n" |
| 54 | |
| 55 | helper.document_get_content = document_get_content |
| 56 | |
| 57 | ok, content = asyncio.run( |
| 58 | helper.document_qa(["/tmp/project.md"], ["What is the codename?"]) |
| 59 | ) |
| 60 | |
| 61 | assert ok is True |
| 62 | assert content == "The project codename is Atlas." |
| 63 | assert "No matching chunks found" in "\n".join(progress) |
| 64 | assert agent.chat_messages is not None |
| 65 | assert "Codename: Atlas" in agent.chat_messages[1].content |
| 66 | |
| 67 | |
| 68 | def test_small_document_fallback_refuses_large_content(): |
| 69 | content = DocumentQueryHelper._small_document_fallback_content( |
| 70 | ["/tmp/large.md"], ["x" * 12_001] |
| 71 | ) |
| 72 | |
| 73 | assert content == "" |