fix: Remove creation of default user during search
This commit is contained in:
parent
38b83a5ec1
commit
a9c507b36e
4 changed files with 3 additions and 19 deletions
|
|
@ -42,14 +42,12 @@ class TripletSearchContextProvider(BaseContextProvider):
|
||||||
self,
|
self,
|
||||||
entities: List[DataPoint],
|
entities: List[DataPoint],
|
||||||
query: str,
|
query: str,
|
||||||
user: User,
|
|
||||||
memory_fragment: CogneeGraph,
|
memory_fragment: CogneeGraph,
|
||||||
) -> List:
|
) -> List:
|
||||||
"""Creates search tasks for valid entities."""
|
"""Creates search tasks for valid entities."""
|
||||||
tasks = [
|
tasks = [
|
||||||
brute_force_triplet_search(
|
brute_force_triplet_search(
|
||||||
query=f"{entity_text} {query}",
|
query=f"{entity_text} {query}",
|
||||||
user=user,
|
|
||||||
top_k=self.top_k,
|
top_k=self.top_k,
|
||||||
collections=self.collections,
|
collections=self.collections,
|
||||||
properties_to_project=self.properties_to_project,
|
properties_to_project=self.properties_to_project,
|
||||||
|
|
@ -84,9 +82,8 @@ class TripletSearchContextProvider(BaseContextProvider):
|
||||||
if not entities:
|
if not entities:
|
||||||
return "No entities provided for context search."
|
return "No entities provided for context search."
|
||||||
|
|
||||||
user = await get_default_user()
|
|
||||||
memory_fragment = await get_memory_fragment(self.properties_to_project)
|
memory_fragment = await get_memory_fragment(self.properties_to_project)
|
||||||
search_tasks = self._get_search_tasks(entities, query, user, memory_fragment)
|
search_tasks = self._get_search_tasks(entities, query, memory_fragment)
|
||||||
|
|
||||||
if not search_tasks:
|
if not search_tasks:
|
||||||
return "No valid entities found for context search."
|
return "No valid entities found for context search."
|
||||||
|
|
|
||||||
|
|
@ -93,11 +93,8 @@ class GraphCompletionRetriever(BaseGraphRetriever):
|
||||||
for field_name in index_fields:
|
for field_name in index_fields:
|
||||||
vector_index_collections.append(f"{subclass.__name__}_{field_name}")
|
vector_index_collections.append(f"{subclass.__name__}_{field_name}")
|
||||||
|
|
||||||
user = await get_default_user()
|
|
||||||
|
|
||||||
found_triplets = await brute_force_triplet_search(
|
found_triplets = await brute_force_triplet_search(
|
||||||
query,
|
query,
|
||||||
user=user,
|
|
||||||
top_k=self.top_k,
|
top_k=self.top_k,
|
||||||
collections=vector_index_collections or None,
|
collections=vector_index_collections or None,
|
||||||
node_type=self.node_type,
|
node_type=self.node_type,
|
||||||
|
|
|
||||||
|
|
@ -89,7 +89,6 @@ async def get_memory_fragment(
|
||||||
|
|
||||||
async def brute_force_triplet_search(
|
async def brute_force_triplet_search(
|
||||||
query: str,
|
query: str,
|
||||||
user: User,
|
|
||||||
top_k: int = 5,
|
top_k: int = 5,
|
||||||
collections: Optional[List[str]] = None,
|
collections: Optional[List[str]] = None,
|
||||||
properties_to_project: Optional[List[str]] = None,
|
properties_to_project: Optional[List[str]] = None,
|
||||||
|
|
@ -102,7 +101,6 @@ async def brute_force_triplet_search(
|
||||||
|
|
||||||
Args:
|
Args:
|
||||||
query (str): The search query.
|
query (str): The search query.
|
||||||
user (User): The user performing the search.
|
|
||||||
top_k (int): The number of top results to retrieve.
|
top_k (int): The number of top results to retrieve.
|
||||||
collections (Optional[List[str]]): List of collections to query.
|
collections (Optional[List[str]]): List of collections to query.
|
||||||
properties_to_project (Optional[List[str]]): List of properties to project.
|
properties_to_project (Optional[List[str]]): List of properties to project.
|
||||||
|
|
@ -139,8 +137,6 @@ async def brute_force_triplet_search(
|
||||||
|
|
||||||
query_vector = (await vector_engine.embedding_engine.embed_text([query]))[0]
|
query_vector = (await vector_engine.embedding_engine.embed_text([query]))[0]
|
||||||
|
|
||||||
send_telemetry("cognee.brute_force_triplet_search EXECUTION STARTED", user.id)
|
|
||||||
|
|
||||||
async def search_in_collection(collection_name: str):
|
async def search_in_collection(collection_name: str):
|
||||||
try:
|
try:
|
||||||
return await vector_engine.search(
|
return await vector_engine.search(
|
||||||
|
|
@ -176,20 +172,15 @@ async def brute_force_triplet_search(
|
||||||
|
|
||||||
results = await memory_fragment.calculate_top_triplet_importances(k=top_k)
|
results = await memory_fragment.calculate_top_triplet_importances(k=top_k)
|
||||||
|
|
||||||
send_telemetry("cognee.brute_force_triplet_search EXECUTION COMPLETED", user.id)
|
|
||||||
|
|
||||||
return results
|
return results
|
||||||
|
|
||||||
except CollectionNotFoundError:
|
except CollectionNotFoundError:
|
||||||
return []
|
return []
|
||||||
except Exception as error:
|
except Exception as error:
|
||||||
logger.error(
|
logger.error(
|
||||||
"Error during brute force search for user: %s, query: %s. Error: %s",
|
"Error during brute force search for query: %s. Error: %s",
|
||||||
user.id,
|
|
||||||
query,
|
query,
|
||||||
error,
|
error,
|
||||||
)
|
)
|
||||||
send_telemetry(
|
send_telemetry("cognee.brute_force_triplet_search EXECUTION FAILED", {"error": str(error)})
|
||||||
"cognee.brute_force_triplet_search EXECUTION FAILED", user.id, {"error": str(error)}
|
|
||||||
)
|
|
||||||
raise error
|
raise error
|
||||||
|
|
|
||||||
|
|
@ -48,7 +48,6 @@ async def main():
|
||||||
query = "When was Kamala Harris in office?"
|
query = "When was Kamala Harris in office?"
|
||||||
triplets = await brute_force_triplet_search(
|
triplets = await brute_force_triplet_search(
|
||||||
query=query,
|
query=query,
|
||||||
user=user,
|
|
||||||
top_k=3,
|
top_k=3,
|
||||||
collections=["graphitinode_content", "graphitinode_name", "graphitinode_summary"],
|
collections=["graphitinode_content", "graphitinode_name", "graphitinode_summary"],
|
||||||
)
|
)
|
||||||
|
|
|
||||||
Loading…
Add table
Reference in a new issue