feat: function categories (#24)

Co-authored-by: afernand <[email protected]>

README.md

@@ -96,6 +96,7 @@ Allie FlowKit Python can be run locally or as a Docker container. Follow the ins
    - Add your function code as an endpoint to a new Python file in the `allie/flowkit/endpoints` directory.
    - Use the `allie/flowkit/endpoints/splitter.py` file and its endpoints as an example.
    - Explicitly define the input and output of the function using Pydantic models, as these will be used by the Allie Agent to call the function.
+   - Add the category and display name of the function to the endpoint definition.
 
 2. **Add the models for the function:**
    - Create the models for the input and output of the function in the `allie/flowkit/models` directory.
@@ -147,7 +148,7 @@ Allie FlowKit Python can be run locally or as a Docker container. Follow the ins
     ```
 
 3. **Define your custom function:**
-- Add your function to ``custom_endpoint.py``, explicitly defining the input and output using Pydantic models.
+- Add your function to ``custom_endpoint.py``, explicitly defining the input and output using Pydantic models, and the category and display name of the function.
 
     **custom_endpoint.py**:
     ```python
@@ -156,6 +157,8 @@ Allie FlowKit Python can be run locally or as a Docker container. Follow the ins
 
 
     @router.post("/custom_function", response_model=CustomResponse)
+    @category(FunctionCategory.GENERIC)
+    @display_name("Custom Function")
     async def custom_function(request: CustomRequest) -> CustomResponse:
         """Endpoint for custom function.
 

src/allie/flowkit/endpoints/splitter.py

@@ -26,7 +26,9 @@
 import io
 
 from allie.flowkit.config._config import CONFIG
+from allie.flowkit.models.functions import FunctionCategory
 from allie.flowkit.models.splitter import SplitterRequest, SplitterResponse
+from allie.flowkit.utils.decorators import category, display_name
 from fastapi import APIRouter, Header, HTTPException
 from langchain.text_splitter import PythonCodeTextSplitter, RecursiveCharacterTextSplitter
 from pdfminer.high_level import extract_text
@@ -38,6 +40,8 @@
 
 
 @router.post("/ppt", response_model=SplitterResponse)
+@category(FunctionCategory.DATA_EXTRACTION)
+@display_name("Split PPT")
 async def split_ppt(request: SplitterRequest, api_key: str = Header(...)) -> SplitterResponse:
     """Endpoint for splitting text in a PowerPoint document into chunks.
 
@@ -55,6 +59,8 @@ async def split_ppt(request: SplitterRequest, api_key: str = Header(...)) -> Spl
 
 
 @router.post("/py", response_model=SplitterResponse)
+@category(FunctionCategory.DATA_EXTRACTION)
+@display_name("Split Python Code")
 async def split_py(request: SplitterRequest, api_key: str = Header(...)) -> SplitterResponse:
     """Endpoint for splitting Python code into chunks.
 
@@ -77,6 +83,8 @@ async def split_py(request: SplitterRequest, api_key: str = Header(...)) -> Spli
 
 
 @router.post("/pdf", response_model=SplitterResponse)
+@category(FunctionCategory.DATA_EXTRACTION)
+@display_name("Split PDF")
 async def split_pdf(request: SplitterRequest, api_key: str = Header(...)) -> SplitterResponse:
     """Endpoint for splitting text in a PDF document into chunks.
 

src/allie/flowkit/fastapi_utils.py

@@ -222,12 +222,22 @@ def extract_endpoint_info(function_map: dict[str, Any], routes: list[APIRoute])
                 output_definitions = get_definitions_from_return_type(return_type) if return_type else {}
                 definitions = {**input_definitions, **output_definitions}
 
+                # Extract category and display name from decorators
+                category = getattr(route.endpoint, "category", "Uncategorized")
+                display_name = getattr(route.endpoint, "display_name", func_name)
+
+                # Extract the description from the docstring
+                description = inspect.getdoc(route.endpoint) or "No description available"
+
                 endpoint_info = EndpointInfo(
                     name=func_name,
                     path=route.path,
                     inputs=inputs,
                     outputs=outputs,
                     definitions=definitions,
+                    category=category,
+                    display_name=display_name,
+                    description=description,
                 )
                 endpoint_list.append(endpoint_info)
     return endpoint_list

src/allie/flowkit/models/functions.py

@@ -22,6 +22,7 @@
 
 """Module for defining the models used in the endpoints."""
 
+from enum import Enum
 from typing import Any
 
 from pydantic import BaseModel
@@ -53,6 +54,19 @@ class EndpointInfo(BaseModel):
 
     name: str
     path: str
+    category: str
+    display_name: str
+    description: str
     inputs: list[ParameterInfo]
     outputs: list[ParameterInfo]
     definitions: dict[str, Any]
+
+
+class FunctionCategory(Enum):
+    """Enum for function categories."""
+
+    DATA_EXTRACTION = "data_extraction"
+    GENERIC = "generic"
+    KNOWLEDGE_DB = "knowledge_db"
+    LLM_HANDLER = "llm_handler"
+    ANSYS_GPT = "ansys_gpt"

src/allie/flowkit/utils/__init__.py

@@ -0,0 +1,22 @@
+# Copyright (C) 2024 ANSYS, Inc. and/or its affiliates.
+# SPDX-License-Identifier: MIT
+#
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+"""Utils module."""

src/allie/flowkit/utils/decorators.py

@@ -0,0 +1,64 @@
+# Copyright (C) 2024 ANSYS, Inc. and/or its affiliates.
+# SPDX-License-Identifier: MIT
+#
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+
+"""Decorators module for function definitions."""
+
+import asyncio
+from functools import wraps
+
+
+def category(value: str):
+    """Decorator to add a category to the function."""
+
+    def decorator(func):
+        func.category = value
+
+        @wraps(func)
+        async def async_wrapper(*args, **kwargs):
+            # Check if function is async
+            if asyncio.iscoroutinefunction(func):
+                return await func(*args, **kwargs)
+            else:
+                return func(*args, **kwargs)
+
+        return async_wrapper
+
+    return decorator
+
+
+def display_name(value: str):
+    """Decorator to add a display name to the function."""
+
+    def decorator(func):
+        func.display_name = value
+
+        @wraps(func)
+        async def async_wrapper(*args, **kwargs):
+            # Check if function is async
+            if asyncio.iscoroutinefunction(func):
+                return await func(*args, **kwargs)
+            else:
+                return func(*args, **kwargs)
+
+        return async_wrapper
+
+    return decorator

tests/test_list_functions.py

@@ -21,6 +21,8 @@
 # SOFTWARE.
 """Test module for list functions."""
 
+import re
+
 from allie.flowkit import flowkit_service
 from fastapi.testclient import TestClient
 import pytest
@@ -29,10 +31,22 @@
 client = TestClient(flowkit_service)
 
 
+def normalize_text(text):
+    """Remove extra spaces, newlines, and indentation."""
+    return re.sub(r"\s+", " ", text.strip())
+
+
+def normalize_response_data(data):
+    """Normalize descriptions in the list of functions."""
+    for item in data:
+        if "description" in item:
+            item["description"] = normalize_text(item["description"])
+    return data
+
+
 @pytest.mark.asyncio
 async def test_list_functions():
     """Test listing available functions."""
-    # Test splitter results
     response = client.get("/", headers={"api-key": "test_api_key"})
     assert response.status_code == 200
     response_data = response.json()
@@ -41,6 +55,16 @@ async def test_list_functions():
         {
             "name": "split_ppt",
             "path": "/splitter/ppt",
+            "category": "data_extraction",
+            "display_name": "Split PPT",
+            "description": """Endpoint for splitting text in a PowerPoint document into chunks.
+            Parameters
+            ----------
+            request : SplitterRequest
+            An object containing 'document_content' in Base64,
+            'chunk_size', and 'chunk_overlap'
+            api_key : str
+            The API key for authentication.""",
             "inputs": [
                 {"name": "document_content", "type": "string(binary)"},
                 {"name": "chunk_size", "type": "integer"},
@@ -52,6 +76,20 @@ async def test_list_functions():
         {
             "name": "split_py",
             "path": "/splitter/py",
+            "category": "data_extraction",
+            "display_name": "Split Python Code",
+            "description": """Endpoint for splitting Python code into chunks.
+            Parameters
+            ----------
+            request : SplitterRequest
+            An object containing 'document_content' in Base64,
+            'chunk_size', and 'chunk_overlap'
+            api_key : str
+            The API key for authentication.
+            Returns
+            -------
+            SplitterResponse
+            An object containing a list of text chunks.""",
             "inputs": [
                 {"name": "document_content", "type": "string(binary)"},
                 {"name": "chunk_size", "type": "integer"},
@@ -63,6 +101,20 @@ async def test_list_functions():
         {
             "name": "split_pdf",
             "path": "/splitter/pdf",
+            "category": "data_extraction",
+            "display_name": "Split PDF",
+            "description": """Endpoint for splitting text in a PDF document into chunks.
+            Parameters
+            ----------
+            request : SplitterRequest
+            An object containing 'document_content' in Base64,
+            'chunk_size', and 'chunk_overlap'.
+            api_key : str
+            The API key for authentication.
+            Returns
+            -------
+            SplitterResponse
+            An object containing a list of text chunks.""",
             "inputs": [
                 {"name": "document_content", "type": "string(binary)"},
                 {"name": "chunk_size", "type": "integer"},
@@ -73,7 +125,11 @@ async def test_list_functions():
         },
     ]
 
-    assert response_data[:3] == expected_response_start
+    # Normalize both actual and expected data
+    normalized_response = normalize_response_data(response_data[:3])
+    normalized_expected = normalize_response_data(expected_response_start)
+
+    assert normalized_response == normalized_expected
 
     # Test invalid API key
     response = client.get("/", headers={"api-key": "invalid_api_key"})