Skip to content

Commit

Permalink
Made response an object with keywords property (#14)
Browse files Browse the repository at this point in the history
This makes it easier to use with the Web API.
  • Loading branch information
de-code authored Dec 19, 2024
1 parent eb35198 commit 5664b74
Show file tree
Hide file tree
Showing 2 changed files with 12 additions and 3 deletions.
11 changes: 9 additions & 2 deletions spacy_keyword_extraction_api/api_router.py
Original file line number Diff line number Diff line change
@@ -1,5 +1,6 @@
import logging
from typing import Sequence
from typing_extensions import TypedDict

from fastapi import APIRouter

Expand All @@ -9,11 +10,17 @@
LOGGER = logging.getLogger(__name__)


class KeywordsResponseTypedDict(TypedDict):
keywords: Sequence[str]


def create_api_router(keyword_extractor: KeywordExtractor) -> APIRouter:
router = APIRouter()

@router.get("/v1/extract-keywords")
def extract_keywords(text: str) -> Sequence[str]:
return list(keyword_extractor.iter_extract_keywords(text_list=[text]))[0]
def extract_keywords(text: str) -> KeywordsResponseTypedDict:
return {
'keywords': list(keyword_extractor.iter_extract_keywords(text_list=[text]))[0]
}

return router
4 changes: 3 additions & 1 deletion tests/unit_test/api_router_test.py
Original file line number Diff line number Diff line change
Expand Up @@ -16,4 +16,6 @@ class TestExtractKeyword:
def test_should_return_keywords_in_lower_case(self):
client = create_test_client()
response = client.get('/v1/extract-keywords', params={'text': 'Some Text'})
assert response.json() == ['some', 'text']
assert response.json() == {
'keywords': ['some', 'text']
}

0 comments on commit 5664b74

Please sign in to comment.