diff --git a/cognitive_architecture/api/v1/search/__init__.py b/cognitive_architecture/api/v1/search/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/cognitive_architecture/api/v1/search/search.py b/cognitive_architecture/api/v1/search/search.py new file mode 100644 index 000000000..5ab34e5c2 --- /dev/null +++ b/cognitive_architecture/api/v1/search/search.py @@ -0,0 +1,33 @@ +""" This module contains the search function that is used to search for nodes in the graph.""" +from enum import Enum, auto +from typing import Dict, Any, Callable, List +from cognitive_architecture.modules.search.graph.search_adjacent import search_adjacent +from cognitive_architecture.modules.search.vector.search_similarity import search_similarity +from cognitive_architecture.modules.search.graph.search_categories import search_categories +from cognitive_architecture.modules.search.graph.search_neighbour import search_neighbour + + +class SearchType(Enum): + ADJACENT = auto() + SIMILARITY = auto() + CATEGORIES = auto() + NEIGHBOR = auto() + + +def complex_search(graph, query_params: Dict[SearchType, Dict[str, Any]]) -> List: + search_functions: Dict[SearchType, Callable] = { + SearchType.ADJACENT: search_adjacent, + SearchType.SIMILARITY: search_similarity, + SearchType.CATEGORIES: search_categories, + SearchType.NEIGHBOR: search_neighbour, + } + + results = set() + + for search_type, params in query_params.items(): + search_func = search_functions.get(search_type) + if search_func: + search_result = search_func(graph, **params) + results.update(search_result) + + return list(results) diff --git a/cognitive_architecture/modules/search/graph/find_proposition_categories.py b/cognitive_architecture/modules/search/graph/find_proposition_categories.py deleted file mode 100644 index 381e4d77e..000000000 --- a/cognitive_architecture/modules/search/graph/find_proposition_categories.py +++ /dev/null @@ -1 +0,0 @@ -#TBD \ No newline at end of file diff --git a/cognitive_architecture/modules/search/graph/find_proposition_neighbourhood.py b/cognitive_architecture/modules/search/graph/find_proposition_neighbourhood.py deleted file mode 100644 index 816e50c28..000000000 --- a/cognitive_architecture/modules/search/graph/find_proposition_neighbourhood.py +++ /dev/null @@ -1,13 +0,0 @@ - -def fetch_context(CONNECTED_GRAPH, id): - relevant_context = [] - for n,attr in CONNECTED_GRAPH.nodes(data=True): - if id in n: - for n_, attr_ in CONNECTED_GRAPH.nodes(data=True): - relevant_layer = attr['layer_uuid'] - - if attr_.get('layer_uuid') == relevant_layer: - print(attr_['description']) - relevant_context.append(attr_['description']) - - return relevant_context diff --git a/cognitive_architecture/modules/search/graph/search_adjacent.py b/cognitive_architecture/modules/search/graph/search_adjacent.py new file mode 100644 index 000000000..aac309488 --- /dev/null +++ b/cognitive_architecture/modules/search/graph/search_adjacent.py @@ -0,0 +1,19 @@ +""" This module contains the function to find the neighbours of a given node in the graph""" + + +def search_adjacent(G, node_id:str)->dict: + """ Find the neighbours of a given node in the graph + :param G: A NetworkX graph object + :param node_id: The unique identifier of the node + :return: A dictionary containing the unique identifiers and descriptions of the neighbours of the given node + """ + + neighbors = list(G.neighbors(node_id)) + neighbor_descriptions = {} + + for neighbor in neighbors: + # Access the 'description' attribute for each neighbor + # The get method returns None if 'description' attribute does not exist for the node + neighbor_descriptions[neighbor] = G.nodes[neighbor].get('description') + + return neighbor_descriptions \ No newline at end of file diff --git a/cognitive_architecture/modules/search/graph/search_categories.py b/cognitive_architecture/modules/search/graph/search_categories.py new file mode 100644 index 000000000..0bf16295b --- /dev/null +++ b/cognitive_architecture/modules/search/graph/search_categories.py @@ -0,0 +1,15 @@ + + + +def search_categories(G, category): + """ + Filter nodes by category. + + Parameters: + - G (nx.Graph): The graph from which to filter nodes. + - category (str): The category to filter nodes by. + + Returns: + - list: A list of nodes that belong to the specified category. + """ + return [node for node, data in G.nodes(data=True) if data.get('category') == category] diff --git a/cognitive_architecture/modules/search/graph/search_neighbour.py b/cognitive_architecture/modules/search/graph/search_neighbour.py new file mode 100644 index 000000000..c800ae5d3 --- /dev/null +++ b/cognitive_architecture/modules/search/graph/search_neighbour.py @@ -0,0 +1,30 @@ +""" Fetches the context of a given node in the graph""" +from cognitive_architecture.infrastructure.databases.graph.get_graph_client import get_graph_client +async def search_neighbour(CONNECTED_GRAPH, id): + relevant_context = [] + for n,attr in CONNECTED_GRAPH.nodes(data=True): + if id in n: + for n_, attr_ in CONNECTED_GRAPH.nodes(data=True): + relevant_layer = attr['layer_uuid'] + + if attr_.get('layer_uuid') == relevant_layer: + print(attr_['description']) + relevant_context.append(attr_['description']) + + return relevant_context + + + +if __name__ == '__main__': + import asyncio + async def main(): + from cognitive_architecture.shared.data_models import GraphDBType + + graph_client = get_graph_client(GraphDBType.NETWORKX) + graph = await graph_client.graph + + await fetch_context(graph, "1") + + asyncio.run(main()) + + diff --git a/cognitive_architecture/modules/search/vector/search_propositions.py b/cognitive_architecture/modules/search/vector/search_propositions.py deleted file mode 100644 index e0b02debf..000000000 --- a/cognitive_architecture/modules/search/vector/search_propositions.py +++ /dev/null @@ -1,15 +0,0 @@ - -async def find_relevant_chunks(query ,unique_layer_uuids): - out = [] - query = await get_embeddings(query) - # print(query) - for id in unique_layer_uuids: - result = qdrant_search(id, query[0]) - - if result: - result_ = [ result_.id for result_ in result] - score_ = [ result_.score for result_ in result] - - out.append([result_, score_]) - - return out \ No newline at end of file diff --git a/cognitive_architecture/modules/search/vector/search_similarity.py b/cognitive_architecture/modules/search/vector/search_similarity.py new file mode 100644 index 000000000..ca78d8666 --- /dev/null +++ b/cognitive_architecture/modules/search/vector/search_similarity.py @@ -0,0 +1,19 @@ + +from cognitive_architecture.infrastructure.llm.get_llm_client import get_llm_client + +async def search_similarity(query ,unique_layer_uuids): + + client = get_llm_client() + out = [] + query = await client.async_get_embedding_with_backoff(query) + # print(query) + for id in unique_layer_uuids: + result = client.search(id, query[0]) + + if result: + result_ = [ result_.id for result_ in result] + score_ = [ result_.score for result_ in result] + + out.append([result_, score_]) + + return out \ No newline at end of file