Add utils for graph visualization + classification nodes

This commit is contained in:
Vasilije 2024-03-12 10:43:22 +01:00
parent 0c9367a2f6
commit afd3053064
9 changed files with 116 additions and 29 deletions

View 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)

View file

@ -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

View file

@ -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

View file

@ -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]

View file

@ -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())

View file

@ -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

View file

@ -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