fix: deletes duplicated retriever instances
This commit is contained in:
hajdul88
parent
2408fd7a01
commit
c094898d15
5 changed files with 0 additions and 169 deletions
|
|
@ -1,25 +0,0 @@
|
||||||
from uuid import UUID
|
|
||||||
from enum import Enum
|
|
||||||
from typing import Callable, Dict
|
|
||||||
from cognee.shared.utils import send_telemetry
|
|
||||||
from cognee.modules.users.models import User
|
|
||||||
from cognee.modules.users.methods import get_default_user
|
|
||||||
from cognee.modules.users.permissions.methods import get_document_ids_for_user
|
|
||||||
|
|
||||||
async def two_step_retriever(query: Dict[str, str], user: User = None) -> list:
|
|
||||||
if user is None:
|
|
||||||
user = await get_default_user()
|
|
||||||
|
|
||||||
if user is None:
|
|
||||||
raise PermissionError("No user found in the system. Please create a user.")
|
|
||||||
|
|
||||||
own_document_ids = await get_document_ids_for_user(user.id)
|
|
||||||
retrieved_results = await diffusion_retriever(query, user)
|
|
||||||
|
|
||||||
filtered_search_results = []
|
|
||||||
|
|
||||||
|
|
||||||
return retrieved_results
|
|
||||||
|
|
||||||
async def diffusion_retriever(query: str, user, community_filter = []) -> list:
|
|
||||||
raise(NotImplementedError)
|
|
||||||
|
|
@ -1,25 +0,0 @@
|
||||||
from uuid import UUID
|
|
||||||
from enum import Enum
|
|
||||||
from typing import Callable, Dict
|
|
||||||
from cognee.shared.utils import send_telemetry
|
|
||||||
from cognee.modules.users.models import User
|
|
||||||
from cognee.modules.users.methods import get_default_user
|
|
||||||
from cognee.modules.users.permissions.methods import get_document_ids_for_user
|
|
||||||
|
|
||||||
async def two_step_retriever(query: Dict[str, str], user: User = None) -> list:
|
|
||||||
if user is None:
|
|
||||||
user = await get_default_user()
|
|
||||||
|
|
||||||
if user is None:
|
|
||||||
raise PermissionError("No user found in the system. Please create a user.")
|
|
||||||
|
|
||||||
own_document_ids = await get_document_ids_for_user(user.id)
|
|
||||||
retrieved_results = await g_retriever(query, user)
|
|
||||||
|
|
||||||
filtered_search_results = []
|
|
||||||
|
|
||||||
|
|
||||||
return retrieved_results
|
|
||||||
|
|
||||||
async def g_retriever(query: str, user, community_filter = []) -> list:
|
|
||||||
raise(NotImplementedError)
|
|
||||||
|
|
@ -1,119 +0,0 @@
|
||||||
import asyncio
|
|
||||||
from uuid import UUID
|
|
||||||
from enum import Enum
|
|
||||||
from typing import Callable, Dict
|
|
||||||
from cognee.shared.utils import send_telemetry
|
|
||||||
from cognee.modules.users.models import User
|
|
||||||
from cognee.modules.users.methods import get_default_user
|
|
||||||
from cognee.modules.users.permissions.methods import get_document_ids_for_user
|
|
||||||
from cognee.modules.graph.cognee_graph.CogneeGraph import CogneeGraph
|
|
||||||
from cognee.infrastructure.databases.vector import get_vector_engine
|
|
||||||
from cognee.infrastructure.databases.graph import get_graph_engine
|
|
||||||
|
|
||||||
|
|
||||||
def format_triplets(edges):
|
|
||||||
print("\n\n\n")
|
|
||||||
def filter_attributes(obj, attributes):
|
|
||||||
"""Helper function to filter out non-None properties, including nested dicts."""
|
|
||||||
result = {}
|
|
||||||
for attr in attributes:
|
|
||||||
value = getattr(obj, attr, None)
|
|
||||||
if value is not None:
|
|
||||||
# If the value is a dict, extract relevant keys from it
|
|
||||||
if isinstance(value, dict):
|
|
||||||
nested_values = {k: v for k, v in value.items() if k in attributes and v is not None}
|
|
||||||
result[attr] = nested_values
|
|
||||||
else:
|
|
||||||
result[attr] = value
|
|
||||||
return result
|
|
||||||
|
|
||||||
triplets = []
|
|
||||||
for edge in edges:
|
|
||||||
node1 = edge.node1
|
|
||||||
node2 = edge.node2
|
|
||||||
edge_attributes = edge.attributes
|
|
||||||
node1_attributes = node1.attributes
|
|
||||||
node2_attributes = node2.attributes
|
|
||||||
|
|
||||||
# Filter only non-None properties
|
|
||||||
node1_info = {key: value for key, value in node1_attributes.items() if value is not None}
|
|
||||||
node2_info = {key: value for key, value in node2_attributes.items() if value is not None}
|
|
||||||
edge_info = {key: value for key, value in edge_attributes.items() if value is not None}
|
|
||||||
|
|
||||||
# Create the formatted triplet
|
|
||||||
triplet = (
|
|
||||||
f"Node1: {node1_info}\n"
|
|
||||||
f"Edge: {edge_info}\n"
|
|
||||||
f"Node2: {node2_info}\n\n\n" # Add three blank lines for separation
|
|
||||||
)
|
|
||||||
triplets.append(triplet)
|
|
||||||
|
|
||||||
return "".join(triplets)
|
|
||||||
|
|
||||||
|
|
||||||
async def two_step_retriever(query: Dict[str, str], user: User = None) -> list:
|
|
||||||
if user is None:
|
|
||||||
user = await get_default_user()
|
|
||||||
|
|
||||||
if user is None:
|
|
||||||
raise PermissionError("No user found in the system. Please create a user.")
|
|
||||||
|
|
||||||
own_document_ids = await get_document_ids_for_user(user.id)
|
|
||||||
retrieved_results = await run_two_step_retriever(query, user)
|
|
||||||
|
|
||||||
filtered_search_results = []
|
|
||||||
|
|
||||||
return retrieved_results
|
|
||||||
|
|
||||||
|
|
||||||
def delete_duplicated_vector_db_elements(collections, results): #:TODO: This is just for now to fix vector db duplicates
|
|
||||||
results_dict = {}
|
|
||||||
for collection, results in zip(collections, results):
|
|
||||||
seen_ids = set()
|
|
||||||
unique_results = []
|
|
||||||
for result in results:
|
|
||||||
if result.id not in seen_ids:
|
|
||||||
unique_results.append(result)
|
|
||||||
seen_ids.add(result.id)
|
|
||||||
else:
|
|
||||||
print(f"Duplicate found in collection '{collection}': {result.id}")
|
|
||||||
results_dict[collection] = unique_results
|
|
||||||
|
|
||||||
return results_dict
|
|
||||||
|
|
||||||
|
|
||||||
async def run_two_step_retriever(query: str, user, community_filter = []) -> list:
|
|
||||||
vector_engine = get_vector_engine()
|
|
||||||
graph_engine = await get_graph_engine()
|
|
||||||
|
|
||||||
collections = ["Entity_name", "TextSummary_text", 'EntityType_name', 'DocumentChunk_text']
|
|
||||||
results = await asyncio.gather(
|
|
||||||
*[vector_engine.get_distances_of_collection(collection, query_text=query) for collection in collections]
|
|
||||||
)
|
|
||||||
|
|
||||||
############################################# This part is a quick fix til we don't fix the vector db inconsistency
|
|
||||||
node_distances = delete_duplicated_vector_db_elements(collections, results)# :TODO: Change when vector db is fixed
|
|
||||||
# results_dict = {collection: result for collection, result in zip(collections, results)}
|
|
||||||
##############################################
|
|
||||||
|
|
||||||
memory_fragment = CogneeGraph()
|
|
||||||
|
|
||||||
await memory_fragment.project_graph_from_db(graph_engine,
|
|
||||||
node_properties_to_project=['id',
|
|
||||||
'description',
|
|
||||||
'name',
|
|
||||||
'type',
|
|
||||||
'text'],
|
|
||||||
edge_properties_to_project=['id',
|
|
||||||
'relationship_name'])
|
|
||||||
|
|
||||||
await memory_fragment.map_vector_distances_to_graph_nodes(node_distances=node_distances)
|
|
||||||
|
|
||||||
await memory_fragment.map_vector_distances_to_graph_edges(vector_engine, query)# :TODO: This should be coming from vector db
|
|
||||||
|
|
||||||
results = await memory_fragment.calculate_top_triplet_importances(k=5)
|
|
||||||
|
|
||||||
print(format_triplets(results))
|
|
||||||
print(f'Query was the following:{query}' )
|
|
||||||
|
|
||||||
return results
|
|
||||||
Loading…
Add table
Reference in a new issue