Add utils for graph visualization + classification nodes
This commit is contained in:
parent
0c9367a2f6
commit
afd3053064
9 changed files with 116 additions and 29 deletions
0
cognitive_architecture/api/v1/search/__init__.py
Normal file
0
cognitive_architecture/api/v1/search/__init__.py
Normal file
33
cognitive_architecture/api/v1/search/search.py
Normal file
33
cognitive_architecture/api/v1/search/search.py
Normal file
|
|
@ -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)
|
||||
|
|
@ -1 +0,0 @@
|
|||
#TBD
|
||||
|
|
@ -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
|
||||
|
|
@ -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
|
||||
|
|
@ -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]
|
||||
|
|
@ -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())
|
||||
|
||||
|
||||
|
|
@ -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
|
||||
|
|
@ -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
|
||||
Loading…
Add table
Reference in a new issue