feat: Added new coding rules search

This commit is contained in:
Igor Ilic 2025-09-04 17:53:07 +02:00
parent 3c50ef4d6f
commit c1106b76fe
5 changed files with 34 additions and 4 deletions

View file

@ -0,0 +1,19 @@
from cognee.shared.logging_utils import get_logger
from cognee.tasks.codingagents.coding_rule_associations import get_existing_rules
logger = get_logger("CodingRulesRetriever")
class CodingRulesRetriever:
"""Retriever for handling codeing rule based searches."""
def __init__(self, rules_nodeset_name):
if isinstance(rules_nodeset_name, list):
rules_nodeset_name = rules_nodeset_name[0]
self.rules_nodeset_name = rules_nodeset_name
"""Initialize retriever with search parameters."""
async def get_existing_rules(self, query_text):
return await get_existing_rules(
rules_nodeset_name=self.rules_nodeset_name, return_list=True
)

View file

@ -13,6 +13,7 @@ from cognee.modules.retrieval.insights_retriever import InsightsRetriever
from cognee.modules.retrieval.summaries_retriever import SummariesRetriever from cognee.modules.retrieval.summaries_retriever import SummariesRetriever
from cognee.modules.retrieval.completion_retriever import CompletionRetriever from cognee.modules.retrieval.completion_retriever import CompletionRetriever
from cognee.modules.retrieval.graph_completion_retriever import GraphCompletionRetriever from cognee.modules.retrieval.graph_completion_retriever import GraphCompletionRetriever
from cognee.modules.retrieval.coding_rules_retriever import CodingRulesRetriever
from cognee.modules.retrieval.graph_summary_completion_retriever import ( from cognee.modules.retrieval.graph_summary_completion_retriever import (
GraphSummaryCompletionRetriever, GraphSummaryCompletionRetriever,
) )
@ -167,6 +168,9 @@ async def specific_search(
SearchType.CYPHER: CypherSearchRetriever().get_completion, SearchType.CYPHER: CypherSearchRetriever().get_completion,
SearchType.NATURAL_LANGUAGE: NaturalLanguageRetriever().get_completion, SearchType.NATURAL_LANGUAGE: NaturalLanguageRetriever().get_completion,
SearchType.FEEDBACK: UserQAFeedback(last_k=last_k).add_feedback, SearchType.FEEDBACK: UserQAFeedback(last_k=last_k).add_feedback,
SearchType.CODING_RULES: CodingRulesRetriever(
rules_nodeset_name=node_name
).get_existing_rules,
} }
# If the query type is FEELING_LUCKY, select the search type intelligently # If the query type is FEELING_LUCKY, select the search type intelligently

View file

@ -15,3 +15,4 @@ class SearchType(Enum):
GRAPH_COMPLETION_CONTEXT_EXTENSION = "GRAPH_COMPLETION_CONTEXT_EXTENSION" GRAPH_COMPLETION_CONTEXT_EXTENSION = "GRAPH_COMPLETION_CONTEXT_EXTENSION"
FEELING_LUCKY = "FEELING_LUCKY" FEELING_LUCKY = "FEELING_LUCKY"
FEEDBACK = "FEEDBACK" FEEDBACK = "FEEDBACK"
CODING_RULES = "CODING_RULES"

View file

@ -31,7 +31,7 @@ class RuleSet(DataPoint):
) )
async def get_existing_rules(rules_nodeset_name: str) -> str: async def get_existing_rules(rules_nodeset_name: str, return_list: bool = False) -> str:
graph_engine = await get_graph_engine() graph_engine = await get_graph_engine()
nodes_data, _ = await graph_engine.get_nodeset_subgraph( nodes_data, _ = await graph_engine.get_nodeset_subgraph(
node_type=NodeSet, node_name=[rules_nodeset_name] node_type=NodeSet, node_name=[rules_nodeset_name]
@ -46,7 +46,8 @@ async def get_existing_rules(rules_nodeset_name: str) -> str:
and "text" in item[1] and "text" in item[1]
] ]
existing_rules = "\n".join(f"- {rule}" for rule in existing_rules) if not return_list:
existing_rules = "\n".join(f"- {rule}" for rule in existing_rules)
return existing_rules return existing_rules

View file

@ -85,8 +85,13 @@ async def main():
) )
# Find the new specific coding rules added to graph through memify (created based on chat conversation between team members) # Find the new specific coding rules added to graph through memify (created based on chat conversation between team members)
developer_rules = await get_existing_rules(rules_nodeset_name="coding_agent_rules") print(
print(developer_rules) await cognee.search(
query_text="List me the coding rules",
query_type=cognee.SearchType.CODING_RULES,
node_name=["coding_agent_rules"],
)
)
# Visualize new graph with added memify context # Visualize new graph with added memify context
file_path = os.path.join( file_path = os.path.join(