fix(auth + brains): handled multiple brains and single brain deployments

Author: ChrisCoder9000

.env.example

@@ -46,5 +46,5 @@ 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"
+# Choose whether to fallback to default brain if not provided
+DEFAULT_BRAIN_FALLBACK="true"

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "brainapi2"
-version = "1.6.6-dev"
+version = "1.6.14-dev"
 description = "Version 1.x.x of the BrainAPI memory layer."
 authors = [
     {name = "Christian",email = "[email protected]"}

src/config.py

@@ -78,7 +78,8 @@ class MilvusConfig:
 
     def __init__(self):
         self.host = os.getenv("MILVUS_HOST")
-        self.port = os.getenv("MILVUS_PORT")
+        port_str = os.getenv("MILVUS_PORT")
+        self.port = int(port_str) if port_str else None
         self.uri = os.getenv("MILVUS_URI")
         self.token = os.getenv("MILVUS_TOKEN")
         if [self.host, self.port].count(None) > 0 and [self.uri, self.token].count(

src/lib/milvus/client.py

@@ -38,6 +38,7 @@ class MilvusClient(VectorStoreClient):
     def __init__(self):
         self._client = None
         self._lock = None
+        self._pid = None
 
     def _get_lock(self):
         if self._lock is None:
@@ -46,23 +47,38 @@ def _get_lock(self):
             self._lock = threading.Lock()
         return self._lock
 
+    def _reset_client_if_forked(self):
+        import os as os_module
+
+        current_pid = os_module.getpid()
+        if self._pid is not None and self._pid != current_pid:
+            if self._client is not None:
+                try:
+                    self._client.close()
+                except Exception:
+                    pass
+            self._client = None
+        self._pid = current_pid
+
     @property
     def client(self):
         """
         Lazy initialization of the Milvus client.
         """
+        self._reset_client_if_forked()
         if self._client is None:
             with self._get_lock():
+                self._reset_client_if_forked()
                 if self._client is None:
                     if config.milvus.uri and config.milvus.token:
                         self._client = Milvus(
                             uri=config.milvus.uri,
                             token=config.milvus.token,
                         )
                     elif config.milvus.host and config.milvus.port:
+                        uri = f"http://{config.milvus.host}:{config.milvus.port}"
                         self._client = Milvus(
-                            host=config.milvus.host,
-                            port=config.milvus.port,
+                            uri=uri,
                             token=config.milvus.token if config.milvus.token else None,
                         )
                     else:

src/lib/neo4j/client.py

@@ -230,6 +230,7 @@ def add_relationship(
         CREATE (a)-[:{":".join(self._clean_labels([predicate.name]))} {{ Description: '{safe_desc}' }}]->(b)
         RETURN a, b
         """
+        self.ensure_database(brain_id)
         result = self.driver.execute_query(cypher_query, database_=brain_id)
         return result
 
@@ -248,6 +249,7 @@ def search_graph(self, nodes: list[Node], brain_id: str) -> list[Node]:
             queries.append(query)
 
         cypher_query = " UNION ".join(queries)
+        self.ensure_database(brain_id)
         result = self.driver.execute_query(cypher_query, database_=brain_id)
         return result
 
@@ -260,6 +262,7 @@ def node_text_search(self, text: str, brain_id: str) -> list[Node]:
         WHERE toLower(n.Name) CONTAINS toLower('{text}')
         RETURN n
         """
+        self.ensure_database(brain_id)
         result = self.driver.execute_query(cypher_query, database_=brain_id)
         return [
             Node(
@@ -315,6 +318,7 @@ def get_nodes_by_uuid(
         if with_relationships:
             cypher_query += ", r, m.uuid as m_uuid, m.name as m_name, labels(m) as m_labels, m.description as m_description, properties(m) as m_properties"
 
+        self.ensure_database(brain_id)
         result = self.driver.execute_query(cypher_query, database_=brain_id)
 
         if with_relationships:
@@ -362,6 +366,7 @@ def get_graph_entities(self, brain_id: str) -> list[str]:
         MATCH (n)
         RETURN DISTINCT labels(n) as labels
         """
+        self.ensure_database(brain_id)
         result = self.driver.execute_query(cypher_query, database_=brain_id)
         return [label for record in result.records for label in record["labels"]]
 
@@ -373,6 +378,7 @@ def get_graph_relationships(self, brain_id: str) -> list[str]:
         CALL db.relationshipTypes() YIELD relationshipType
         RETURN relationshipType
         """
+        self.ensure_database(brain_id)
         result = self.driver.execute_query(cypher_query, database_=brain_id)
         return [record["relationshipType"] for record in result.records]
 
@@ -384,6 +390,7 @@ def get_by_uuid(self, uuid: str, brain_id: str) -> Node:
         MATCH (n) WHERE n.uuid = '{uuid}'
         RETURN n.uuid as uuid, n.name as name, labels(n) as labels, n.description as description, properties(n) as properties
         """
+        self.ensure_database(brain_id)
         result = self.driver.execute_query(cypher_query, database_=brain_id)
         if not result.records or len(result.records) == 0:
             return None
@@ -403,6 +410,7 @@ def get_by_uuids(self, uuids: list[str], brain_id: str) -> list[Node]:
         MATCH (n) WHERE n.uuid IN ["{'","'.join(uuids)}"]
         RETURN n.uuid as uuid, n.name as name, labels(n) as labels, n.description as description, properties(n) as properties
         """
+        self.ensure_database(brain_id)
         result = self.driver.execute_query(cypher_query, database_=brain_id)
         return [
             Node(
@@ -438,6 +446,7 @@ def get_by_identification_params(
         MATCH (n{(":" + ":".join(self._clean_labels(entity_types))) if entity_types else ""}) {("WHERE " + " AND ".join(where_clauses)) if where_clauses else ""}
         RETURN n.uuid as uuid, n.name as name, labels(n) as labels, n.description as description, properties(n) as properties
         """
+        self.ensure_database(brain_id)
         result = self.driver.execute_query(cypher_query, database_=brain_id)
         if not result.records or len(result.records) == 0:
             return None
@@ -457,6 +466,7 @@ def get_graph_property_keys(self, brain_id: str) -> list[str]:
         CALL db.propertyKeys() YIELD propertyKey
         RETURN propertyKey
         """
+        self.ensure_database(brain_id)
         result = self.driver.execute_query(cypher_query, database_=brain_id)
         return [record["propertyKey"] for record in result.records]
 
@@ -474,6 +484,7 @@ def get_neighbors(
                CASE WHEN startNode(r2) = c THEN 'out' ELSE 'in' END AS direction,
                c.uuid AS c_uuid, c.name AS c_name, labels(c) AS c_labels, c.description AS c_description, properties(c) AS c_properties
         """
+        self.ensure_database(brain_id)
         result = self.driver.execute_query(cypher_query, database_=brain_id)
 
         neighbors = []
@@ -538,6 +549,7 @@ def get_neighbor_node_tuples(
             m.uuid as m_uuid, m.name as m_name, labels(m) as m_labels,
             m.description as m_description, properties(m) as m_properties, r as rel
         """
+        self.ensure_database(brain_id)
         result = self.driver.execute_query(cypher_query, database_=brain_id)
 
         if not result.records:
@@ -632,6 +644,7 @@ def get_connected_nodes(
             n.uuid as n_uuid, n.name as n_name, labels(n) as n_labels, n.description as n_description, properties(n) as n_properties,
             CASE WHEN startNode(r) = n THEN 'out' ELSE 'in' END AS direction
         """
+        self.ensure_database(brain_id)
         result = self.driver.execute_query(cypher_query, database_=brain_id)
         return [
             (
@@ -727,6 +740,7 @@ def search_relationships(
         {"WHERE " + " AND ".join(filters) if filters else ""}
         RETURN count(r) AS total
         """
+        self.ensure_database(brain_id)
         result = self.driver.execute_query(cypher_query, database_=brain_id)
         count_result = self.driver.execute_query(cypher_count, database_=brain_id)
         total = 0
@@ -854,6 +868,7 @@ def search_entities(
         {"WHERE " + " AND ".join(filters) if filters else ""}
         RETURN count(n) AS total
         """
+        self.ensure_database(brain_id)
         result = self.driver.execute_query(cypher_query, database_=brain_id)
         count_result = self.driver.execute_query(cypher_count, database_=brain_id)
         total = 0

src/services/api/constants/requests.py

@@ -155,6 +155,9 @@ class RetrieveNeighborsAiModeRequestBody(BaseModel):
         ...,
         description="The description of the neighbors to look for.",
     )
+    brain_id: str = Field(
+        default="default", description="The brain identifier to store the data in."
+    )
 
 
 class RetrieveNeighborsWithIdentificationParamsRequestBody(BaseModel):
@@ -167,3 +170,14 @@ class RetrieveNeighborsWithIdentificationParamsRequestBody(BaseModel):
         description="The identification parameters of the entity to get neighbors for.",
     )
     limit: int = Field(10, description="The number of neighbors to return.")
+    brain_id: str = Field(
+        default="default", description="The brain identifier to store the data in."
+    )
+
+
+class CreateBrainRequest(BaseModel):
+    """
+    Request body for the create brain endpoint.
+    """
+
+    brain_id: str

src/services/api/controllers/system.py

@@ -11,6 +11,7 @@
 import asyncio
 
 from pydantic import BaseModel
+from src.services.api.constants.requests import CreateBrainRequest
 from src.services.data.main import data_adapter
 
 
@@ -22,14 +23,6 @@ async def get_brains_list():
     return result
 
 
-class CreateBrainRequest(BaseModel):
-    """
-    Request body for the create brain endpoint.
-    """
-
-    brain_id: str
-
-
 async def create_new_brain(request: CreateBrainRequest):
     """
     Create a new brain

src/services/api/middlewares/auth.py

@@ -23,51 +23,50 @@ async def dispatch(self, request: Request, call_next):
         if request.method == "OPTIONS":
             return await call_next(request)
 
+        # Variables ----------------------------------------------
         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'}"
+        system_pat = os.getenv("BRAINPAT_TOKEN")
 
-        use_only_system_pat = os.getenv("USE_ONLY_SYSTEM_PAT") == "true"
-
-        if use_only_system_pat:
-            cachepat_key = "brainpat:system"
+        if request.url.path.startswith("/system") or request.url.path == "/":
+            if brainpat == system_pat:
+                return await call_next(request)
+            return JSONResponse(
+                status_code=status.HTTP_401_UNAUTHORIZED,
+                content={"detail": "Invalid or missing BrainPAT header"},
+            )
 
+        brain_id = getattr(request.state, "brain_id", None)
+        if not brain_id:
+            return JSONResponse(
+                status_code=status.HTTP_400_BAD_REQUEST,
+                content={"detail": "Brain ID is required."},
+            )
+        cachepat_key = f"brainpat:{brain_id}"
         cached_brainpat = cache_adapter.get(key=cachepat_key, brain_id="system")
 
-        if not cached_brainpat and not use_only_system_pat:
+        # Logic --------------------------------------------------
+        if brainpat == system_pat:
+            return await call_next(request)
+
+        if not cached_brainpat:
             stored_brain = data_adapter.get_brain(name_key=brain_id)
-            system_pat = os.getenv("BRAINPAT_TOKEN")
-            if not stored_brain:
-                if brainpat != system_pat:
-                    return JSONResponse(
-                        status_code=status.HTTP_401_UNAUTHORIZED,
-                        content={"detail": "Invalid or missing BrainPAT header"},
-                    )
-                cached_brainpat = system_pat
-            else:
-                cached_brainpat = stored_brain.pat
-                cache_adapter.set(
-                    key=cachepat_key,
-                    value=stored_brain.pat,
-                    brain_id="system",
+            if not stored_brain or stored_brain.pat != brainpat:
+                return JSONResponse(
+                    status_code=status.HTTP_401_UNAUTHORIZED,
+                    content={"detail": "Invalid or missing BrainPAT header"},
                 )
-        if not cached_brainpat and use_only_system_pat:
-            system_pat = os.getenv("BRAINPAT_TOKEN")
-            cached_brainpat = system_pat
+            cached_brainpat = stored_brain.pat
             cache_adapter.set(
-                key="brainpat:system",
-                value=system_pat,
-                brain_id="system",
+                key=cachepat_key, value=stored_brain.pat, brain_id="system"
             )
-
-        if cached_brainpat != brainpat:
+        elif 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