[WIP] Allow to implement custom endpoints

core/cat/looking_glass/cheshire_cat.py

@@ -2,6 +2,9 @@
 from typing import List, Dict
 from typing_extensions import Protocol
 
+from fastapi.routing import APIRouter
+from fastapi import FastAPI
+
 from langchain_core.language_models.llms import BaseLLM
 from langchain.base_language import BaseLanguageModel
 from langchain_core.language_models.chat_models import BaseChatModel
@@ -82,6 +85,11 @@ def __init__(self):
         # allows plugins to do something after the cat bootstrap is complete
         self.mad_hatter.execute_hook("after_cat_bootstrap", cat=self)
 
+    def load_custom_endpoints(self):
+        custom_endpoints_router = APIRouter(prefix="/custom")
+        custom_endpoints_router = self.mad_hatter.execute_hook("load_custom_endpoints", custom_endpoints_router, cat=self)
+        return custom_endpoints_router
+
     def load_natural_language(self):
         """Load Natural Language related objects.
 

core/cat/mad_hatter/core_plugin/hooks/flow.py

@@ -6,6 +6,7 @@
 
 from cat.mad_hatter.decorators import hook
 from langchain.docstore.document import Document
+from fastapi.routing import APIRouter
 
 
 # Called before cat bootstrap
@@ -331,4 +332,26 @@ def before_cat_stores_episodic_memory(doc: Document, cat) -> Document:
             `when`: timestamp to track when it's been uploaded.
 
     """
-    return doc
+    return doc
+
+
+# Hook called to extend custom endpoint routes in cat rest api
+@hook(priority=0)
+def load_custom_endpoints(custom_endpoint_router: APIRouter, cat) -> None:
+    """Hook the custom endpoints router loading.
+
+    Allow to extend custom endpoint routes in FastAPI
+
+    Parameters
+    ----------
+    custom_endpoint_router : APIRouter
+        FastAPI APIRouter instance.
+    cat : CheshireCat
+        Cheshire Cat instance.
+
+    Returns
+    -------
+    custom_endpoint_router : APIRouter
+        extended FastAPI APIRouter instance with custom routes.
+    """
+    return custom_endpoint_router

core/cat/main.py

@@ -28,6 +28,9 @@ async def lifespan(app: FastAPI):
     # - Starlette allows this: https://www.starlette.io/applications/#storing-state-on-the-app-instance
     app.state.ccat = CheshireCat()
 
+    # Load custom endpoints from cat tools
+    app.include_router(router=app.state.ccat.load_custom_endpoints())
+
     # Dict of pseudo-sessions (key is the user_id)
     app.state.strays = {}