Skip to content

Commit

Permalink
feat: Add support for retrieving EngineFields
Browse files Browse the repository at this point in the history
  • Loading branch information
Milan Ondrašovič committed Jan 14, 2025
1 parent 3ba3e90 commit 00b86c0
Show file tree
Hide file tree
Showing 6 changed files with 140 additions and 7 deletions.
3 changes: 2 additions & 1 deletion rossum_api/domain_logic/resources.py
Original file line number Diff line number Diff line change
Expand Up @@ -14,6 +14,8 @@ class Resource(Enum):
Connector = "connectors"
Document = "documents"
EmailTemplate = "email_templates"
Engine = "engines"
EngineField = "engine_fields"
Group = "groups"
Hook = "hooks"
Inbox = "inboxes"
Expand All @@ -24,4 +26,3 @@ class Resource(Enum):
Upload = "uploads"
User = "users"
Workspace = "workspaces"
Engine = "engines"
23 changes: 22 additions & 1 deletion rossum_api/elis_api_client.py
Original file line number Diff line number Diff line change
Expand Up @@ -24,7 +24,7 @@
from rossum_api.models.connector import Connector
from rossum_api.models.document import Document
from rossum_api.models.email_template import EmailTemplate
from rossum_api.models.engine import Engine
from rossum_api.models.engine import Engine, EngineField
from rossum_api.models.group import Group
from rossum_api.models.hook import Hook
from rossum_api.models.inbox import Inbox
Expand Down Expand Up @@ -561,6 +561,27 @@ async def retrieve_engine(self, engine_id: int) -> Engine:

return self._deserializer(Resource.Engine, engine)

async def list_all_engines(
self,
ordering: Sequence[str] = (),
sideloads: Sequence[str] = (),
**filters: Any,
) -> AsyncIterator[Engine]:
"""https://elis.rossum.ai/api/docs/internal/#list-all-engines."""
async for engine in self._http_client.fetch_all(
Resource.Engine, ordering, sideloads, **filters
):
yield self._deserializer(Resource.Engine, engine)

async def retrieve_engine_fields(
self, engine_id: int | None = None
) -> AsyncIterator[EngineField]:
"""https://elis.rossum.ai/api/docs/internal/#engine-field."""
async for engine_field in self._http_client.fetch_all(
Resource.EngineField, engine=engine_id
):
yield self._deserializer(Resource.EngineField, engine_field)

# ##### INBOX #####
async def create_new_inbox(self, data: Dict[str, Any]) -> Inbox:
"""https://elis.rossum.ai/api/docs/#create-a-new-inbox."""
Expand Down
17 changes: 16 additions & 1 deletion rossum_api/elis_api_client_sync.py
Original file line number Diff line number Diff line change
Expand Up @@ -34,7 +34,7 @@
from rossum_api.models.connector import Connector
from rossum_api.models.document import Document
from rossum_api.models.email_template import EmailTemplate
from rossum_api.models.engine import Engine
from rossum_api.models.engine import Engine, EngineField
from rossum_api.models.group import Group
from rossum_api.models.hook import Hook
from rossum_api.models.inbox import Inbox
Expand Down Expand Up @@ -283,6 +283,21 @@ def retrieve_engine(self, engine_id: int) -> Engine:
"""https://elis.rossum.ai/api/docs/#retrieve-a-schema."""
return self._run_coroutine(self.elis_api_client.retrieve_engine(engine_id))

async def list_all_engines(
self,
ordering: Sequence[str] = (),
sideloads: Sequence[str] = (),
**filters: Any,
) -> AsyncIterator[Engine]:
"""https://elis.rossum.ai/api/docs/internal/#list-all-engines."""
return self._run_coroutine(
self.elis_api_client.list_all_engines(ordering, sideloads, **filters)
)

def retrieve_engine_fields(self, engine_id: int | None = None) -> list[EngineField]:
"""https://elis.rossum.ai/api/docs/internal/#engine-field."""
return self._run_coroutine(self.elis_api_client.retrieve_engine_fields(engine_id))

# ##### USERS #####
def list_all_users(
self,
Expand Down
5 changes: 3 additions & 2 deletions rossum_api/models/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -10,7 +10,7 @@
from rossum_api.models.connector import Connector
from rossum_api.models.document import Document
from rossum_api.models.email_template import EmailTemplate
from rossum_api.models.engine import Engine
from rossum_api.models.engine import Engine, EngineField
from rossum_api.models.group import Group
from rossum_api.models.hook import Hook
from rossum_api.models.inbox import Inbox
Expand All @@ -34,6 +34,8 @@
Resource.Connector: Connector,
Resource.Document: Document,
Resource.EmailTemplate: EmailTemplate,
Resource.Engine: Engine,
Resource.EngineField: EngineField,
Resource.Group: Group,
Resource.Hook: Hook,
Resource.Inbox: Inbox,
Expand All @@ -44,7 +46,6 @@
Resource.Upload: Upload,
Resource.User: User,
Resource.Workspace: Workspace,
Resource.Engine: Engine,
}


Expand Down
24 changes: 22 additions & 2 deletions rossum_api/models/engine.py
Original file line number Diff line number Diff line change
@@ -1,14 +1,34 @@
from __future__ import annotations

from dataclasses import dataclass
from dataclasses import dataclass, field
from typing import Literal, Optional

EngineFieldType = Literal["string", "number", "date", "enum", "button"]
MultilineType = Literal["true", "false", ""] # Preparation for "auto" option


@dataclass
class Engine:
id: int
url: str
name: str
type: str
# Workout around the issue that "type" is a reserved keyword
type: str = field(metadata={"dataclasses_json": {"name": "type"}})
learning_enabled: bool
description: str
agenda_id: str


@dataclass
class EngineField:
id: int
url: str
engine: str
name: str
tabular: bool
label: str
# Workout around the issue that "type" is a reserved keyword
type: EngineFieldType = field(metadata={"dataclasses_json": {"name": "type"}})
subtype: Optional[str]
pre_trained_field_id: Optional[str]
multiline: MultilineType
75 changes: 75 additions & 0 deletions tests/elis_api_client/test_engines.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,75 @@
from __future__ import annotations

import pytest

from rossum_api.domain_logic.resources import Resource
from rossum_api.models.engine import Engine, EngineField

TEST_ENGINE_ID = 123


@pytest.mark.asyncio
class TestEngine:
@pytest.fixture
def dummy_engine(self):
# https://elis.rossum.ai/api/docs/internal/#engine
return {
"id": TEST_ENGINE_ID,
"url": f"https://elis.rossum.ai/api/v1/engines/{TEST_ENGINE_ID}",
"name": "test_engine",
"type": "extractor",
"learning_enabled": False,
"description": "Test engine",
"agenda_id": "test_agenda_id",
}

async def test_retrieve_engine(self, elis_client, dummy_engine):
client, http_client = elis_client
http_client.fetch_one.return_value = dummy_engine

engine: Engine = await client.retrieve_engine(TEST_ENGINE_ID)

assert engine == Engine(**dummy_engine)
http_client.fetch_one.assert_called_with(Resource.Engine, TEST_ENGINE_ID)

async def test_list_all_engines(self, elis_client, dummy_engine, mock_generator):
client, http_client = elis_client
http_client.fetch_all.return_value = mock_generator(dummy_engine)

engines = client.list_all_engines()

async for engine in engines:
assert engine == Engine(**dummy_engine)

http_client.fetch_all.assert_called_with(Resource.Engine, (), ())


@pytest.mark.asyncio
class TestEngineFields:
@pytest.fixture
def dummy_engine_field(self):
# https://elis.rossum.ai/api/docs/internal/#engine-field
return {
"id": 456, # EngineField ID, not the Engine ID
"url": f"https://elis.rossum.ai/api/v1/engine_fields/{TEST_ENGINE_ID}",
"engine": f"https://elis.rossum.ai/api/v1/engines/{TEST_ENGINE_ID}",
"name": f"test_engine_field_{TEST_ENGINE_ID}",
"label": "Test engine field",
"type": "string",
"subtype": "alphanumeric",
"pre_trained_field_id": "document_id", # https://elis.rossum.ai/api/docs/internal/#get-list-of-possible-pre_trained_field_id-fields
"tabular": False,
"multiline": "false",
}

async def test_retrieve_engine_fields(self, elis_client, dummy_engine_field, mock_generator):
client, http_client = elis_client

http_client.fetch_all.return_value = mock_generator(dummy_engine_field)

engine_fields = client.retrieve_engine_fields(TEST_ENGINE_ID)

async for engine_field in engine_fields:
assert engine_field == EngineField(**dummy_engine_field)

http_client.fetch_all.assert_called_with(Resource.EngineField, engine=TEST_ENGINE_ID)

0 comments on commit 00b86c0

Please sign in to comment.