feat: Added new coding rules search
This commit is contained in:
parent
3c50ef4d6f
commit
c1106b76fe
5 changed files with 34 additions and 4 deletions
19
cognee/modules/retrieval/coding_rules_retriever.py
Normal file
19
cognee/modules/retrieval/coding_rules_retriever.py
Normal 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
|
||||||
|
)
|
||||||
|
|
@ -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
|
||||||
|
|
|
||||||
|
|
@ -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"
|
||||||
|
|
|
||||||
|
|
@ -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
|
||||||
|
|
||||||
|
|
|
||||||
|
|
@ -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(
|
||||||
|
|
|
||||||
Loading…
Add table
Reference in a new issue