feat(pat): multiple pat x brain

Author: ChrisCoder9000

.env.example

@@ -42,3 +42,9 @@ MONGO_PASSWORD="password"
 # Auth
 # Used to handle authentication
 BRAINPAT_TOKEN="your_token"
+
+# MultiBrain
+# Choose to allow or block automatic creation of new brains on requests with non existing new brain_ids
+BRAIN_CREATION_ALLOWED="true"
+# Choose whether to use for every brain the main pat or to use dedicated one for each brain
+USE_ONLY_SYSTEM_PAT="false"

src/constants/data.py

@@ -90,3 +90,17 @@ class Brain(BaseModel):
 
     id: str = Field(default_factory=lambda: str(uuid.uuid4()))
     name_key: str = Field(description="The key used to identify the brain.")
+
+    @staticmethod
+    def _random_pat() -> str:
+        import random
+
+        chars = []
+        for _ in range(48):
+            chars.append(random.choice("abcdefghijklmnopqrstuvwxyz0123456789"))
+        return "".join(chars)
+
+    pat: str = Field(
+        description="The personal access token for the brain.",
+        default_factory=_random_pat,
+    )

src/lib/mongo/client.py

@@ -134,21 +134,32 @@ def save_structured_data(
 
     def create_brain(self, name_key: str) -> Brain:
         collection = self.get_collection("brains", "system")
-        result = collection.insert_one({"name_key": name_key})
-        return Brain(id=str(result.inserted_id), name_key=name_key)
+        brain = Brain(name_key=name_key)
+        brain_dict = brain.model_dump(mode="json", exclude={"id"})
+        result = collection.insert_one(brain_dict)
+        brain.id = str(result.inserted_id)
+        return brain
 
     def get_brain(self, name_key: str) -> Brain:
         collection = self.get_collection("brains", "system")
         result = collection.find_one({"name_key": name_key})
         if not result:
             return None
-        return Brain(id=str(result["_id"]), name_key=result["name_key"])
+        return Brain(
+            id=str(result["_id"]),
+            name_key=result["name_key"],
+            pat=result.get("pat"),
+        )
 
     def get_brains_list(self) -> List[Brain]:
         collection = self.get_collection("brains", "system")
         result = collection.find()
         return [
-            Brain(id=str(result["_id"]), name_key=result["name_key"])
+            Brain(
+                id=str(result["_id"]),
+                name_key=result["name_key"],
+                pat=result.get("pat"),
+            )
             for result in result
         ]
 

src/lib/redis/client.py

@@ -43,7 +43,12 @@ def get(self, key: str, brain_id: str) -> str:
         Get a value from the cache.
         """
         prefixed_key = self._get_key(key, brain_id)
-        return self.client.get(prefixed_key)
+        result = self.client.get(prefixed_key)
+        if result is None:
+            return None
+        if isinstance(result, bytes):
+            return result.decode("utf-8")
+        return result
 
     def set(
         self, key: str, value: str, brain_id: str, expires_in: Optional[int] = None

src/services/api/app.py

@@ -25,14 +25,15 @@
 app = FastAPI()
 
 app.add_middleware(BrainPATMiddleware)
+app.add_middleware(BrainMiddleware)
 app.add_middleware(
     CORSMiddleware,
     allow_origins=["*"],
     allow_credentials=False,
     allow_methods=["*"],
     allow_headers=["*"],
 )
-app.add_middleware(BrainMiddleware)
+
 
 app.include_router(ingest_router)
 app.include_router(retrieve_router)

src/services/api/controllers/system.py

@@ -18,3 +18,11 @@ async def get_brains_list():
     """
     result = await asyncio.to_thread(data_adapter.get_brains_list)
     return result
+
+
+async def create_new_brain(brain_id: str):
+    """
+    Create a new brain
+    """
+    result = await asyncio.to_thread(data_adapter.create_brain, brain_id)
+    return result

src/services/api/middlewares/auth.py

@@ -8,23 +8,64 @@
 -----
 """
 
+import os
 from fastapi import Request
 from fastapi.responses import JSONResponse
 from starlette.middleware.base import BaseHTTPMiddleware
 from starlette import status
 
-from src.config import config
+from src.services.kg_agent.main import cache_adapter
+from src.services.data.main import data_adapter
 
 
 class BrainPATMiddleware(BaseHTTPMiddleware):
     async def dispatch(self, request: Request, call_next):
         if request.method == "OPTIONS":
             return await call_next(request)
-        brainpat = request.headers.get("BrainPAT")
-        if brainpat != config.brainpat_token:
+
+        brainpat = request.headers.get("BrainPAT") or getattr(
+            request.state, "pat", None
+        )
+        brain_id = getattr(request.state, "brain_id", None)
+
+        cachepat_key = f"brainpat:{brain_id or 'default'}"
+
+        use_only_system_pat = os.getenv("USE_ONLY_SYSTEM_PAT") == "true"
+
+        if use_only_system_pat:
+            cachepat_key = "brainpat:system"
+
+        cached_brainpat = cache_adapter.get(key=cachepat_key, brain_id="system")
+
+        if not cached_brainpat and not use_only_system_pat:
+            stored_brain = data_adapter.get_brain(name_key=brain_id)
+            system_pat = os.getenv("BRAINPAT_TOKEN")
+            if not stored_brain and brainpat != system_pat:
+                return JSONResponse(
+                    status_code=status.HTTP_401_UNAUTHORIZED,
+                    content={"detail": "Invalid or missing BrainPAT header"},
+                )
+
+            cached_brainpat = stored_brain.pat
+            cache_adapter.set(
+                key=cachepat_key,
+                value=stored_brain.pat,
+                brain_id="system",
+            )
+        if not cached_brainpat and use_only_system_pat:
+            system_pat = os.getenv("BRAINPAT_TOKEN")
+            cached_brainpat = system_pat
+            cache_adapter.set(
+                key="brainpat:system",
+                value=system_pat,
+                brain_id="system",
+            )
+
+        if cached_brainpat != brainpat:
             return JSONResponse(
                 status_code=status.HTTP_401_UNAUTHORIZED,
                 content={"detail": "Invalid or missing BrainPAT header"},
             )
         response = await call_next(request)
+
         return response

src/services/api/middlewares/brains.py

@@ -8,6 +8,7 @@
 -----
 """
 
+import os
 from fastapi import Request
 from fastapi.responses import JSONResponse
 from starlette.middleware.base import BaseHTTPMiddleware
@@ -62,9 +63,11 @@ async def receive():
                 key=f"brain:{brain_id}", brain_id="system"
             )
 
-            if not cached_brain_id:
-                stored_brain = data_adapter.get_brain(name_key=brain_id)
+            brain_creation_allowed = os.getenv("BRAIN_CREATION_ALLOWED") == "true"
 
+            if not cached_brain_id and brain_creation_allowed:
+                stored_brain = data_adapter.get_brain(name_key=brain_id)
+                new_brain = None
                 if not stored_brain:
                     new_brain = data_adapter.create_brain(name_key=brain_id)
                     cache_adapter.set(
@@ -79,6 +82,18 @@ async def receive():
                         brain_id="system",
                     )
 
+                cached_brain_id = new_brain.id if new_brain else stored_brain.id
+
+                use_only_system_pat = os.getenv("USE_ONLY_SYSTEM_PAT") == "true"
+                if not use_only_system_pat:
+                    request.state.pat = new_brain.pat if new_brain else stored_brain.pat
+
+            if not cached_brain_id:
+                return JSONResponse(
+                    status_code=status.HTTP_406_NOT_ACCEPTABLE,
+                    content={"detail": "Brain not found or creation is not allowed."},
+                )
+
         request.state.brain_id = brain_id
 
         return await call_next(request)

src/services/api/routes/system.py

@@ -11,6 +11,7 @@
 from fastapi import APIRouter
 from src.services.api.controllers.system import (
     get_brains_list as get_brains_list_controller,
+    create_new_brain as create_new_brain_controller,
 )
 
 system_router = APIRouter(prefix="/system", tags=["system"])
@@ -22,3 +23,11 @@ async def get_brains_list():
     Get the list of brains.
     """
     return await get_brains_list_controller()
+
+
+@system_router.post(path="/brains")
+async def create_brain():
+    """
+    Create a new brain
+    """
+    return await create_new_brain_controller()