113 lines
4.6 KiB
Python
113 lines
4.6 KiB
Python
import os
|
|
import tempfile
|
|
import asyncio
|
|
import unittest
|
|
from pathlib import Path
|
|
from unittest.mock import AsyncMock, patch
|
|
|
|
_TEST_DATA_DIR = None
|
|
if "SOVEREIGNAI_DATA_DIR" not in os.environ:
|
|
_TEST_DATA_DIR = tempfile.TemporaryDirectory(prefix="sovereignai-mixed-pdf-index-tests-")
|
|
os.environ["SOVEREIGNAI_DATA_DIR"] = _TEST_DATA_DIR.name
|
|
|
|
from app import database, knowledge
|
|
from app.main import app
|
|
from tests.api_client import authenticated_client
|
|
from tests.test_pdf_analysis import make_mixed_pdf
|
|
|
|
|
|
class MixedPdfKnowledgeIntegrationTests(unittest.TestCase):
|
|
def test_indexes_and_searches_digital_and_scanned_pages_together(self) -> None:
|
|
data_dir = Path(os.environ["SOVEREIGNAI_DATA_DIR"])
|
|
data_dir.mkdir(parents=True, exist_ok=True)
|
|
workspace = data_dir / "mixed-pdf-project"
|
|
workspace.mkdir(exist_ok=True)
|
|
pdf_path = workspace / "mixed.pdf"
|
|
pdf_path.write_bytes(make_mixed_pdf())
|
|
database_path = data_dir / "mixed-pdf-knowledge.sqlite3"
|
|
ocr = {
|
|
"engine": "easyocr-local-ar-en",
|
|
"text": "ScannedIndexMarker ArabicPageText",
|
|
"lines": [],
|
|
"average_confidence": 0.85,
|
|
"elapsed_seconds": 0.1,
|
|
}
|
|
|
|
async def fake_embeddings(texts, *, model=None):
|
|
return [[1.0, 0.0] for _ in texts]
|
|
|
|
try:
|
|
with (
|
|
patch.object(database, "DATABASE_PATH", database_path),
|
|
patch.dict(
|
|
os.environ,
|
|
{
|
|
"KNOWLEDGE_EMBEDDING_MODEL": "granite-embedding:278m",
|
|
"SOVEREIGNAI_ALLOWED_WORKSPACES": str(data_dir),
|
|
},
|
|
),
|
|
patch("app.main.embeddings.embed_texts", new=AsyncMock(side_effect=fake_embeddings)),
|
|
):
|
|
database.initialize_database()
|
|
knowledge.initialize()
|
|
client = authenticated_client(app)
|
|
ocr_execution_contexts: list[str] = []
|
|
|
|
def fake_ocr(*_args, **_kwargs):
|
|
try:
|
|
asyncio.get_running_loop()
|
|
except RuntimeError:
|
|
ocr_execution_contexts.append("worker")
|
|
else:
|
|
ocr_execution_contexts.append("event-loop")
|
|
return ocr
|
|
|
|
with patch("app.main.recognize_image_text", side_effect=fake_ocr):
|
|
indexed = client.post(
|
|
"/v1/agent/knowledge/index",
|
|
json={"workspace_path": str(workspace), "files": ["mixed.pdf"]},
|
|
)
|
|
digital = client.post(
|
|
"/v1/agent/knowledge/search",
|
|
json={"workspace_path": str(workspace), "task": "IndexMarker"},
|
|
)
|
|
scanned = client.post(
|
|
"/v1/agent/knowledge/search",
|
|
json={"workspace_path": str(workspace), "task": "ScannedIndexMarker"},
|
|
)
|
|
semantic = client.post(
|
|
"/v1/agent/knowledge/search",
|
|
json={"workspace_path": str(workspace), "task": "صياغة مختلفة بلا كلمات مشتركة"},
|
|
)
|
|
removed = client.request(
|
|
"DELETE",
|
|
"/v1/agent/knowledge/index",
|
|
json={"workspace_path": str(workspace), "files": ["mixed.pdf"]},
|
|
)
|
|
|
|
self.assertEqual(indexed.status_code, 200, indexed.text)
|
|
self.assertEqual(indexed.json()["indexed"][0]["ocr_pages"], 1)
|
|
self.assertEqual(ocr_execution_contexts, ["worker"])
|
|
self.assertTrue(indexed.json()["indexed"][0]["semantic_indexed"])
|
|
self.assertEqual(digital.status_code, 200, digital.text)
|
|
self.assertIn("IndexMarker", digital.json()["results"][0]["text"])
|
|
self.assertEqual(scanned.status_code, 200, scanned.text)
|
|
self.assertIn("ScannedIndexMarker", scanned.json()["results"][0]["text"])
|
|
self.assertEqual(semantic.json()["search_mode"], "hybrid")
|
|
self.assertTrue(semantic.json()["results"])
|
|
self.assertEqual(removed.status_code, 200, removed.text)
|
|
self.assertTrue(removed.json()["deleted"][0]["deleted"])
|
|
self.assertFalse(
|
|
knowledge.has_embeddings(
|
|
user_id=database.LOCAL_USER_ID,
|
|
workspace_path=workspace,
|
|
model="granite-embedding:278m",
|
|
)
|
|
)
|
|
finally:
|
|
pdf_path.unlink(missing_ok=True)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|