d27f847753
<!-- .github/pull_request_template.md --> ## Description Delete legacy search implementations after migrating to new retriever classes ## DCO Affirmation I affirm that all code in every commit of this pull request conforms to the terms of the Topoteretes Developer Certificate of Origin <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit - **New Features** - Enhanced search and retrieval capabilities, providing improved context resolution for code queries, completions, summaries, and graph connections. - **Refactor** - Shifted to a modular, object-oriented approach that consolidates query logic and streamlines error management for a more robust and scalable experience. - **Bug Fixes** - Improved error handling for unsupported search types and retrieval operations. <!-- end of auto-generated comment: release notes by coderabbit.ai -->
80 lines
2.5 KiB
Python
80 lines
2.5 KiB
Python
import asyncio
|
||
|
||
import cognee
|
||
import logging
|
||
from cognee.modules.pipelines import Task, run_tasks
|
||
from cognee.shared.utils import setup_logging
|
||
from cognee.tasks.temporal_awareness import build_graph_with_temporal_awareness
|
||
from cognee.infrastructure.databases.relational import (
|
||
create_db_and_tables as create_relational_db_and_tables,
|
||
)
|
||
from cognee.tasks.temporal_awareness.index_graphiti_objects import (
|
||
index_and_transform_graphiti_nodes_and_edges,
|
||
)
|
||
from cognee.modules.retrieval.utils.brute_force_triplet_search import brute_force_triplet_search
|
||
from cognee.modules.retrieval.graph_completion_retriever import GraphCompletionRetriever
|
||
from cognee.infrastructure.llm.prompts import read_query_prompt, render_prompt
|
||
from cognee.infrastructure.llm.get_llm_client import get_llm_client
|
||
|
||
text_list = [
|
||
"Kamala Harris is the Attorney General of California. She was previously "
|
||
"the district attorney for San Francisco.",
|
||
"As AG, Harris was in office from January 3, 2011 – January 3, 2017",
|
||
]
|
||
|
||
|
||
async def main():
|
||
await cognee.prune.prune_data()
|
||
await cognee.prune.prune_system(metadata=True)
|
||
await create_relational_db_and_tables()
|
||
|
||
for text in text_list:
|
||
await cognee.add(text)
|
||
|
||
tasks = [
|
||
Task(build_graph_with_temporal_awareness, text_list=text_list),
|
||
]
|
||
|
||
pipeline = run_tasks(tasks)
|
||
|
||
async for result in pipeline:
|
||
print(result)
|
||
|
||
await index_and_transform_graphiti_nodes_and_edges()
|
||
|
||
query = "When was Kamala Harris in office?"
|
||
triplets = await brute_force_triplet_search(
|
||
query=query,
|
||
top_k=3,
|
||
collections=["graphitinode_content", "graphitinode_name", "graphitinode_summary"],
|
||
)
|
||
|
||
retriever = GraphCompletionRetriever()
|
||
context = await retriever.resolve_edges_to_text(triplets)
|
||
|
||
args = {
|
||
"question": query,
|
||
"context": context,
|
||
}
|
||
|
||
user_prompt = render_prompt("graph_context_for_question.txt", args)
|
||
system_prompt = read_query_prompt("answer_simple_question_restricted.txt")
|
||
|
||
llm_client = get_llm_client()
|
||
computed_answer = await llm_client.acreate_structured_output(
|
||
text_input=user_prompt,
|
||
system_prompt=system_prompt,
|
||
response_model=str,
|
||
)
|
||
|
||
print(computed_answer)
|
||
|
||
|
||
if __name__ == "__main__":
|
||
setup_logging(logging.ERROR)
|
||
loop = asyncio.new_event_loop()
|
||
asyncio.set_event_loop(loop)
|
||
try:
|
||
loop.run_until_complete(main())
|
||
finally:
|
||
loop.run_until_complete(loop.shutdown_asyncgens())
|