feat(brains)

Author: ChrisCoder9000

src/adapters/cache.py

@@ -8,6 +8,7 @@
 -----
 """
 
+from typing import Optional
 from .interfaces.cache import CacheClient
 
 
@@ -25,20 +26,26 @@ def add_client(self, client: CacheClient) -> None:
         """
         self.cache = client
 
-    def get(self, key: str) -> str:
+    def get(self, key: str, brain_id: str = "default") -> str:
         """
         Get a value from the cache.
         """
-        return self.cache.get(key)
+        return self.cache.get(key, brain_id)
 
-    def set(self, key: str, value: str, expires_in: int) -> bool:
+    def set(
+        self,
+        key: str,
+        value: str,
+        brain_id: str = "default",
+        expires_in: Optional[int] = None,
+    ) -> bool:
         """
         Set a value in the cache with an expiration time.
         """
-        return self.cache.set(key, value, expires_in)
+        return self.cache.set(key, value, brain_id, expires_in)
 
-    def delete(self, key: str) -> bool:
+    def delete(self, key: str, brain_id: str = "default") -> bool:
         """
         Delete a value from the cache.
         """
-        return self.cache.delete(key)
+        return self.cache.delete(key, brain_id)

src/adapters/data.py

@@ -10,7 +10,7 @@
 
 from typing import List, Tuple
 from src.adapters.interfaces.data import DataClient, SearchResult
-from src.constants.data import Observation, StructuredData, TextChunk
+from src.constants.data import Brain, Observation, StructuredData, TextChunk
 
 
 class DataAdapter:
@@ -27,34 +27,52 @@ def add_client(self, client: DataClient) -> None:
         """
         self.data = client
 
-    def save_text_chunk(self, text_chunk: TextChunk) -> TextChunk:
+    def save_text_chunk(
+        self, text_chunk: TextChunk, brain_id: str = "default"
+    ) -> TextChunk:
         """
         Save a text chunk to the data client.
         """
-        return self.data.save_text_chunk(text_chunk)
+        return self.data.save_text_chunk(text_chunk, brain_id)
 
-    def save_observations(self, observations: List[Observation]) -> Observation:
+    def save_observations(
+        self, observations: List[Observation], brain_id: str = "default"
+    ) -> Observation:
         """
         Save a list of observations to the data client.
         """
-        return self.data.save_observations(observations)
+        return self.data.save_observations(observations, brain_id)
 
-    def search(self, text: str) -> SearchResult:
+    def search(self, text: str, brain_id: str = "default") -> SearchResult:
         """
         Search data by text and return a list of text chunks and observations.
         """
-        return self.data.search(text)
+        return self.data.search(text, brain_id)
 
     def get_text_chunks_by_ids(
-        self, ids: List[str], with_observations: bool
+        self, ids: List[str], with_observations: bool, brain_id: str = "default"
     ) -> Tuple[List[TextChunk], List[Observation]]:
         """
         Get data by their IDs.
         """
-        return self.data.get_text_chunks_by_ids(ids, with_observations)
+        return self.data.get_text_chunks_by_ids(ids, with_observations, brain_id)
 
-    def save_structured_data(self, structured_data: StructuredData) -> StructuredData:
+    def save_structured_data(
+        self, structured_data: StructuredData, brain_id: str = "default"
+    ) -> StructuredData:
         """
         Save a structured data to the data client.
         """
-        return self.data.save_structured_data(structured_data)
+        return self.data.save_structured_data(structured_data, brain_id)
+
+    def create_brain(self, name_key: str) -> Brain:
+        """
+        Create a new brain in the data client.
+        """
+        return self.data.create_brain(name_key)
+
+    def get_brain(self, name_key: str) -> Brain:
+        """
+        Get a brain from the data client.
+        """
+        return self.data.get_brain(name_key)

src/adapters/embeddings.py

@@ -41,32 +41,43 @@ def add_client(self, client: VectorStoreClient) -> None:
         """
         self.vector_store = client
 
-    def add_vectors(self, vectors: list[Vector], store: str) -> None:
+    def add_vectors(
+        self, vectors: list[Vector], store: str, brain_id: str = "default"
+    ) -> None:
         """
         Add vectors to the vector store.
         """
-        return self.vector_store.add_vectors(vectors, store)
+        return self.vector_store.add_vectors(vectors, store, brain_id)
 
-    def search_vectors(self, query: str, store: str, k: int = 10) -> list[Vector]:
+    def search_vectors(
+        self, query: str, store: str, brain_id: str = "default", k: int = 10
+    ) -> list[Vector]:
         """
         Search vectors in the vector store and return the top k vectors.
         """
-        return self.vector_store.search_vectors(query, store, k)
+        return self.vector_store.search_vectors(query, store, brain_id, k)
 
-    def get_by_ids(self, ids: list[str], store: str) -> list[Vector]:
+    def get_by_ids(
+        self, ids: list[str], store: str, brain_id: str = "default"
+    ) -> list[Vector]:
         """
         Get vectors by their IDs.
         """
-        return self.vector_store.get_by_ids(ids, store)
+        return self.vector_store.get_by_ids(ids, store, brain_id)
 
     def search_similar_by_ids(
-        self, vector_ids: list[str], store: str, min_similarity: float, limit: int = 10
+        self,
+        vector_ids: list[str],
+        store: str,
+        min_similarity: float,
+        limit: int = 10,
+        brain_id: str = "default",
     ) -> list[Vector]:
         """
         Search similar vectors by their IDs.
         """
         return self.vector_store.search_similar_by_ids(
-            vector_ids, store, min_similarity, limit
+            vector_ids, store, min_similarity, limit, brain_id
         )
 
 

src/adapters/graph.py

@@ -16,7 +16,6 @@
     Predicate,
     SearchEntitiesResult,
     SearchRelationshipsResult,
-    Triple,
 )
 
 
@@ -50,57 +49,60 @@ def add_client(self, client: GraphClient) -> None:
         """
         self.graph = client
 
-    def execute_operation(self, operation: str) -> str:
+    def execute_operation(self, operation: str, brain_id: str = "default") -> str:
         """
         Execute a generic graph operation.
         """
         try:
-            return self.graph.execute_operation(operation)
+            return self.graph.execute_operation(operation, brain_id)
         except Exception as e:  # pylint: disable=broad-exception-caught
             print(f"Error executing graph operation: {e} - {operation}")
             return f"Error executing graph operation: {e}"
 
     def add_nodes(
         self,
         nodes: list[Node],
+        brain_id: str = "default",
         identification_params: Optional[dict] = None,
         metadata: Optional[dict] = None,
-        database: Optional[str] = None,
     ) -> list[Node] | str:
         """
         Add nodes to the graph.
         """
-        return self.graph.add_nodes(nodes, identification_params, metadata, database)
+        return self.graph.add_nodes(nodes, brain_id, identification_params, metadata)
 
     def add_relationship(
         self,
         subject: Node,
         predicate: Predicate,
         to_object: Node,
+        brain_id: str = "default",
     ) -> str:
         """
         Add a relationship between two nodes to the graph.
         """
-        return self.graph.add_relationship(subject, predicate, to_object)
+        return self.graph.add_relationship(subject, predicate, to_object, brain_id)
 
     def search_graph(
         self,
         nodes: list[Node],
+        brain_id: str = "default",
     ) -> list[Node]:
         """
         Search the graph for nodes and 1 degree relationships.
         """
-        return self.graph.search_graph(nodes)
+        return self.graph.search_graph(nodes, brain_id)
 
-    def node_text_search(self, text: str) -> list[Node]:
+    def node_text_search(self, text: str, brain_id: str = "default") -> list[Node]:
         """
         Search the graph for nodes by partial text match into the name of the nodes.
         """
-        return self.graph.node_text_search(text)
+        return self.graph.node_text_search(text, brain_id)
 
     def get_nodes_by_uuid(
         self,
         uuids: list[str],
+        brain_id: str = "default",
         with_relationships: Optional[bool] = False,
         relationships_depth: Optional[int] = 1,
         relationships_type: Optional[list[str]] = None,
@@ -111,80 +113,83 @@ def get_nodes_by_uuid(
         """
         return self.graph.get_nodes_by_uuid(
             uuids,
+            brain_id,
             with_relationships,
             relationships_depth,
             relationships_type,
             preferred_labels,
         )
 
-    def get_graph_entities(self) -> list[str]:
+    def get_graph_entities(self, brain_id: str = "default") -> list[str]:
         """
         Get the entities of the graph.
         """
-        return self.graph.get_graph_entities()
+        return self.graph.get_graph_entities(brain_id)
 
-    def get_graph_relationships(self) -> list[str]:
+    def get_graph_relationships(self, brain_id: str = "default") -> list[str]:
         """
         Get the relationships of the graph.
         """
-        return self.graph.get_graph_relationships()
+        return self.graph.get_graph_relationships(brain_id)
 
-    def get_graph_property_keys(self) -> list[str]:
+    def get_graph_property_keys(self, brain_id: str = "default") -> list[str]:
         """
         Get the property keys of the graph.
         """
-        return self.graph.get_graph_property_keys()
+        return self.graph.get_graph_property_keys(brain_id)
 
-    def get_by_uuid(self, uuid: str) -> Node:
+    def get_by_uuid(self, uuid: str, brain_id: str = "default") -> Node:
         """
         Get a node by its UUID.
         """
-        return self.graph.get_by_uuid(uuid)
+        return self.graph.get_by_uuid(uuid, brain_id)
 
-    def get_by_uuids(self, uuids: list[str]) -> list[Node]:
+    def get_by_uuids(self, uuids: list[str], brain_id: str = "default") -> list[Node]:
         """
         Get nodes by their UUIDs.
         """
-        return self.graph.get_by_uuids(uuids)
+        return self.graph.get_by_uuids(uuids, brain_id)
 
     def get_by_identification_params(
         self,
         identification_params: IdentificationParams,
+        brain_id: str = "default",
         entity_types: Optional[list[str]] = None,
     ) -> Node:
         """
         Get a node by its identification params and entity types.
         """
         return self.graph.get_by_identification_params(
-            identification_params, entity_types
+            identification_params, brain_id, entity_types
         )
 
     def get_neighbors(
-        self, node: Node, limit: int
+        self, node: Node, limit: int, brain_id: str = "default"
     ) -> list[Tuple[Node, Predicate, Node]]:
         """
         Get the neighbors of a node.
         """
-        return self.graph.get_neighbors(node, limit)
+        return self.graph.get_neighbors(node, limit, brain_id)
 
     def get_node_with_rel_by_uuid(
-        self, rel_ids_with_node_ids: list[tuple[str, str]]
+        self, rel_ids_with_node_ids: list[tuple[str, str]], brain_id: str = "default"
     ) -> list[dict]:
         """
         Get the node with the relationships by their UUIDs.
         """
-        return self.graph.get_node_with_rel_by_uuid(rel_ids_with_node_ids)
+        return self.graph.get_node_with_rel_by_uuid(rel_ids_with_node_ids, brain_id)
 
     def get_neighbor_node_tuples(
-        self, a_uuid: str, b_uuids: list[str]
+        self, a_uuid: str, b_uuids: list[str], brain_id: str = "default"
     ) -> list[Tuple[Node, Predicate, Node]]:
         """
         Get the neighbor node tuples by their UUIDs.
         """
-        return self.graph.get_neighbor_node_tuples(a_uuid, b_uuids)
+        return self.graph.get_neighbor_node_tuples(a_uuid, b_uuids, brain_id)
 
     def get_connected_nodes(
         self,
+        brain_id: str = "default",
         node: Optional[Node] = None,
         uuids: Optional[list[str]] = None,
         limit: Optional[int] = 10,
@@ -194,11 +199,12 @@ def get_connected_nodes(
         Get the connected nodes by their UUIDs.
         """
         return self.graph.get_connected_nodes(
-            node=node, uuids=uuids, limit=limit, with_labels=with_labels
+            brain_id, node=node, uuids=uuids, limit=limit, with_labels=with_labels
         )
 
     def search_relationships(
         self,
+        brain_id: str = "default",
         limit: int = 10,
         skip: int = 0,
         relationship_types: Optional[list[str]] = None,
@@ -213,6 +219,7 @@ def search_relationships(
         relationship_uuids = []
         # TODO: semantic search + src/core/agents/tools/kg_agent/KGAgentAddTripletsTool.py:165
         return self.graph.search_relationships(
+            brain_id,
             limit,
             skip,
             relationship_types,
@@ -225,6 +232,7 @@ def search_relationships(
 
     def search_entities(
         self,
+        brain_id: str = "default",
         limit: int = 10,
         skip: int = 0,
         node_labels: Optional[list[str]] = None,
@@ -236,7 +244,7 @@ def search_entities(
         node_uuids = []
         # TODO: semantic search
         return self.graph.search_entities(
-            limit, skip, node_labels, node_uuids, query_text
+            brain_id, limit, skip, node_labels, node_uuids, query_text
         )