feat: add definitions endpoint, minor refactoring

Signed-off-by: Panos Vagenas <35837085+vagenas@users.noreply.github.com>

deepsearch/cli.py

@@ -4,7 +4,6 @@
 from deepsearch.core.cli.plugins import get_cli_groups
 from deepsearch.cps.cli.main import app as cps_app
 from deepsearch.documents.cli.main import app as documents_app
-from deepsearch.model.cli.main import app as model_app
 from deepsearch.query.cli.main import app as query_app
 
 app.add_typer(cps_app, name="cps", help="Interact with DeepSearch CPS component")
@@ -14,9 +13,6 @@
     name="documents",
     help="Interact with DeepSearch Document Conversion component",
 )
-app.add_typer(
-    model_app, name="model", help="Interact with the DeepSearch model component"
-)
 
 for group in get_cli_groups():
     app.add_typer(group)

deepsearch/model/README.md

@@ -0,0 +1,24 @@
+# Model API
+
+> Currently in **beta**.
+
+The Model API allows users to serve and integrate their own models.
+
+## Usage
+```python
+from deepsearch.model.server.deepsearch_annotator_app import DeepSearchAnnotatorApp
+
+annotator = ...  # e.g. SimpleTextGeographyAnnotator()
+app = DeepSearchAnnotatorApp()
+app.register_annotator(annotator)
+app.run()
+```
+
+For a complete example, check [examples/main.py](examples/main.py).
+
+## Models
+For an example, check
+[examples/simple_text_geography_annotator/](examples/simple_text_geography_annotator/).
+
+## Base models
+Check [base/](base/).

deepsearch/model/examples/main.py

@@ -0,0 +1,15 @@
+from deepsearch.model.examples.simple_text_geography_annotator.simple_text_geography_annotator import (  # type: ignore
+    SimpleTextGeographyAnnotator,
+)
+from deepsearch.model.server.deepsearch_annotator_app import DeepSearchAnnotatorApp
+
+
+def run():
+    app = DeepSearchAnnotatorApp()
+    app.register_annotator(SimpleTextGeographyAnnotator())
+    app.register_annotator(SimpleTextGeographyAnnotator(), name="foo")
+    app.run()
+
+
+if __name__ == "__main__":
+    run()

deepsearch/model/server/deepsearch_annotator_app.py

@@ -4,12 +4,11 @@
 import inspect
 import logging
 import os
-import sys
 import time
-from concurrent.futures import ThreadPoolExecutor
 from datetime import datetime
-from typing import Coroutine, Dict, List, Optional, Union
+from typing import Coroutine, Dict, Union
 
+import uvicorn
 from anyio import CapacityLimiter
 from anyio.lowlevel import RunVar
 from fastapi import FastAPI, HTTPException, Request
@@ -66,9 +65,12 @@ async def validation_exception_handler(
         async def health_check() -> dict:
             return {"message": "HealthCheck"}
 
-        @self.app.get("/annotator")
-        async def health_check() -> dict:
-            return {key: f"annotator/{key}" for key in self.annotators_list}
+        @self.app.get("/")
+        async def get_definitions() -> dict:
+            return {
+                key: self.annotate_controller.get_annotator_info(key)
+                for key in self.annotators_list
+            }
 
         # Will Require an API key
         @self.app.get("/annotator/{annotator_name}")
@@ -197,6 +199,9 @@ def register_annotator(
             )
             exit(-1)
 
+    def run(self, host: str = "127.0.0.1", port: int = 8000, **kwargs) -> None:
+        uvicorn.run(self.app, host=host, port=port, **kwargs)
+
     class AnnotateController:
         def __init__(self, app_annotators):
             self.app_annotators: dict = app_annotators