mirror of
https://github.com/agent0ai/agent-zero.git
synced 2026-05-17 04:01:13 +00:00
72 lines
1.9 KiB
Python
72 lines
1.9 KiB
Python
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import sys
|
|
from pathlib import Path
|
|
|
|
PROJECT_ROOT = Path(__file__).resolve().parents[1]
|
|
if str(PROJECT_ROOT) not in sys.path:
|
|
sys.path.insert(0, str(PROJECT_ROOT))
|
|
|
|
from helpers.document_query import DocumentQueryHelper
|
|
|
|
|
|
class FakeStore:
|
|
@staticmethod
|
|
def normalize_uri(uri: str) -> str:
|
|
return uri
|
|
|
|
async def search_documents(self, **_kwargs):
|
|
return []
|
|
|
|
|
|
class FakeAgent:
|
|
def __init__(self):
|
|
self.chat_messages = None
|
|
|
|
async def handle_intervention(self):
|
|
return None
|
|
|
|
def parse_prompt(self, name: str) -> str:
|
|
return name
|
|
|
|
async def call_utility_model(self, **_kwargs) -> str:
|
|
return "codename"
|
|
|
|
async def call_chat_model(self, messages, explicit_caching=False):
|
|
self.chat_messages = messages
|
|
return "The project codename is Atlas.", None
|
|
|
|
|
|
def test_document_qa_uses_small_document_content_when_search_finds_no_chunks():
|
|
agent = FakeAgent()
|
|
progress = []
|
|
helper = object.__new__(DocumentQueryHelper)
|
|
helper.agent = agent
|
|
helper.store = FakeStore()
|
|
helper.progress_callback = progress.append
|
|
|
|
async def document_get_content(uri, add_to_db=False):
|
|
assert uri == "/tmp/project.md"
|
|
assert add_to_db is True
|
|
return "# Project\n\nCodename: Atlas\n"
|
|
|
|
helper.document_get_content = document_get_content
|
|
|
|
ok, content = asyncio.run(
|
|
helper.document_qa(["/tmp/project.md"], ["What is the codename?"])
|
|
)
|
|
|
|
assert ok is True
|
|
assert content == "The project codename is Atlas."
|
|
assert "No matching chunks found" in "\n".join(progress)
|
|
assert agent.chat_messages is not None
|
|
assert "Codename: Atlas" in agent.chat_messages[1].content
|
|
|
|
|
|
def test_small_document_fallback_refuses_large_content():
|
|
content = DocumentQueryHelper._small_document_fallback_content(
|
|
["/tmp/large.md"], ["x" * 12_001]
|
|
)
|
|
|
|
assert content == ""
|