pip install cognee[codegraph]

import os
os.environ["LLM_API_KEY"] = "sk-your_actual_api_key_here"  # Replace with your actual API key

git clone https://github.com/hande-k/simple-repo.git

repo_path = "/path/to/your/simple-repo"  # Adjust this path to your cloned repo location

import cognee
from cognee.api.v1.cognify.code_graph_pipeline import run_code_graph_pipeline

async def codify(repo_path: str):
    """Run the code graph pipeline on the specified repository"""
    print("\nStarting code graph pipeline...")
    async for result in run_code_graph_pipeline(repo_path, False):
        print(result)
    print("\nPipeline completed!")

await codify(repo_path)

with open("summarize_search_results.txt", "w") as f:
    f.write(
        "You are a helpful assistant that understands the given user query "
        "and the results returned based on the query. Provide a concise, "
        "short, to-the-point user-friendly explanation based on these."
    )

from cognee.modules.search.types import SearchType
from cognee.infrastructure.llm.prompts import read_query_prompt
from cognee.infrastructure.llm.get_llm_client import get_llm_client

async def retrieve_and_generate_answer(query: str) -> str:
    """Search the code graph and generate a human-friendly answer"""
    
    # Search the code graph
    search_results = await cognee.search(
        query_type=SearchType.CODE, 
        query_text=query
    )
    
    # Load the summarization prompt
    prompt_path = "summarize_search_results.txt"  # Adjust path if needed
    system_prompt = read_query_prompt(prompt_path)
    
    # Get LLM client
    llm_client = get_llm_client()

    # Generate summary
    answer = await llm_client.acreate_structured_output(
        text_input=(
            f"Search Results:\n{str(search_results)}\n\n"
            f"User Query:\n{query}\n"
        ),
        system_prompt=system_prompt,
        response_model=str,
    )
    
    return answer

# Example queries - replace with your own
user_queries = [
    "What functions are available in this codebase?",
    "How does the main application work?",
    "What are the key classes and their relationships?",
    "Show me the data flow in this application"
]

for query in user_queries:
    print(f"\n🤔 Query: {query}")
    answer = await retrieve_and_generate_answer(query)
    print("📋 Answer:")
    print(answer)
    print("-" * 50)

# Ask your own question
user_query = "How is user authentication handled in this codebase?"
answer = await retrieve_and_generate_answer(user_query)
print("===== ANSWER =====")
print(answer)

# Get raw code chunks
chunks = await cognee.search(
    query_type=SearchType.CHUNKS, 
    query_text="authentication logic"
)

# Get insights about relationships
insights = await cognee.search(
    query_type=SearchType.INSIGHTS, 
    query_text="how modules interact"
)