expose doc_item_refs and labels in sandbox search results
This commit is contained in:
parent
4118533db1
commit
bacc21b38b
3 changed files with 110 additions and 0 deletions
|
|
@ -65,6 +65,8 @@ class Sandbox:
|
||||||
"score": r.score,
|
"score": r.score,
|
||||||
"page_numbers": r.page_numbers,
|
"page_numbers": r.page_numbers,
|
||||||
"headings": r.headings,
|
"headings": r.headings,
|
||||||
|
"doc_item_refs": r.doc_item_refs,
|
||||||
|
"labels": r.labels,
|
||||||
}
|
}
|
||||||
for r in expanded
|
for r in expanded
|
||||||
]
|
]
|
||||||
|
|
|
||||||
|
|
@ -132,6 +132,32 @@ class TestSandboxHaikuRAG:
|
||||||
assert result.success
|
assert result.success
|
||||||
assert "True" in result.stdout or "1" in result.stdout
|
assert "True" in result.stdout or "1" in result.stdout
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
@pytest.mark.vcr()
|
||||||
|
async def test_search_returns_doc_item_refs_and_labels(self, temp_db_path):
|
||||||
|
"""Search results include doc_item_refs and labels."""
|
||||||
|
config = AppConfig()
|
||||||
|
async with HaikuRAG(temp_db_path, create=True) as client:
|
||||||
|
await client.create_document(
|
||||||
|
content="The quick brown fox jumps over the lazy dog.",
|
||||||
|
uri="test://animals",
|
||||||
|
title="Animals",
|
||||||
|
)
|
||||||
|
|
||||||
|
context = AnalysisContext()
|
||||||
|
sb = Sandbox(client=client, config=config, context=context)
|
||||||
|
result = await sb.execute(
|
||||||
|
"results = await search('fox', limit=1)\n"
|
||||||
|
"r = results[0]\n"
|
||||||
|
"print('doc_item_refs' in r)\n"
|
||||||
|
"print('labels' in r)\n"
|
||||||
|
"print(type(r['doc_item_refs']).__name__)\n"
|
||||||
|
"print(type(r['labels']).__name__)"
|
||||||
|
)
|
||||||
|
assert result.success
|
||||||
|
assert "True\nTrue" in result.stdout
|
||||||
|
assert "list\nlist" in result.stdout
|
||||||
|
|
||||||
@pytest.mark.asyncio
|
@pytest.mark.asyncio
|
||||||
@pytest.mark.vcr()
|
@pytest.mark.vcr()
|
||||||
async def test_get_document(self, temp_db_path):
|
async def test_get_document(self, temp_db_path):
|
||||||
|
|
|
||||||
File diff suppressed because one or more lines are too long
Loading…
Reference in a new issue