refactor: Refactor search so graph completion is used by default (#505)

<!-- .github/pull_request_template.md -->

## Description
Refactor search so query type doesn't need to be provided to make it
simpler for new users

## 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

- **Refactor**
- Improved the search interface by standardizing parameter usage with
explicit keyword arguments for specifying search types, enhancing
clarity and consistency.
- **Tests**
- Updated test cases and example integrations to align with the revised
search parameters, ensuring consistent behavior and reliable validation
of search outcomes.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->

.github/workflows/test_dynamic_steps_example.yml

@@ -16,13 +16,7 @@ jobs:
       with:
         example-location: ./examples/python/dynamic_steps_example.py
       secrets:
-        LLM_MODEL: ${{ secrets.LLM_MODEL }}
+        LLM_API_KEY: ${{ secrets.OPENAI_API_KEY }}
-        LLM_ENDPOINT: ${{ secrets.LLM_ENDPOINT }}
+        OPENAI_API_KEY: ${{ secrets.OPENAI_API_KEY }}
-        LLM_API_KEY: ${{ secrets.LLM_API_KEY }}
-        LLM_API_VERSION: ${{ secrets.LLM_API_VERSION }}
-        EMBEDDING_MODEL: ${{ secrets.EMBEDDING_MODEL }}
-        EMBEDDING_ENDPOINT: ${{ secrets.EMBEDDING_ENDPOINT }}
-        EMBEDDING_API_KEY: ${{ secrets.EMBEDDING_API_KEY }}
-        EMBEDDING_API_VERSION: ${{ secrets.EMBEDDING_API_VERSION }}
         GRAPHISTRY_USERNAME: ${{ secrets.GRAPHISTRY_USERNAME }}
         GRAPHISTRY_PASSWORD: ${{ secrets.GRAPHISTRY_PASSWORD }}

.github/workflows/test_llama_index_cognee_integration_notebook.yml

@@ -9,20 +9,51 @@ concurrency:
   group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }}
   cancel-in-progress: true
 
+env:
+  RUNTIME__LOG_LEVEL: ERROR
+
 jobs:
   run_notebook_test:
-      uses: ./.github/workflows/reusable_notebook.yml
+    name: test
-      with:
+    runs-on: ubuntu-22.04
-        notebook-location: notebooks/llama_index_cognee_integration.ipynb
+    defaults:
-      secrets:
+      run:
-        #LLM_MODEL: ${{ secrets.LLM_MODEL }}
+        shell: bash
-        #LLM_ENDPOINT: ${{ secrets.LLM_ENDPOINT }}
+    steps:
-        LLM_API_KEY: ${{ secrets.OPENAI_API_KEY }}
+      - name: Check out
-        OPENAI_API_KEY: ${{ secrets.OPENAI_API_KEY }}
+        uses: actions/checkout@master
-        #LLM_API_VERSION: ${{ secrets.LLM_API_VERSION }}
+
-        EMBEDDING_MODEL: ${{ secrets.EMBEDDING_MODEL }}
+      - name: Setup Python
-        EMBEDDING_ENDPOINT: ${{ secrets.EMBEDDING_ENDPOINT }}
+        uses: actions/setup-python@v5
-        EMBEDDING_API_KEY: ${{ secrets.EMBEDDING_API_KEY }}
+        with:
-        EMBEDDING_API_VERSION: ${{ secrets.EMBEDDING_API_VERSION }}
+          python-version: '3.11.x'
-        GRAPHISTRY_USERNAME: ${{ secrets.GRAPHISTRY_USERNAME }}
+
-        GRAPHISTRY_PASSWORD: ${{ secrets.GRAPHISTRY_PASSWORD }}
+      - name: Install Poetry
+        uses: snok/install-poetry@v1.4.1
+        with:
+          virtualenvs-create: true
+          virtualenvs-in-project: true
+          installer-parallel: true
+
+      - name: Install dependencies
+        run: |
+          pip install jupyter
+          pip install llama-index-graph-rag-cognee==0.1.2
+
+      - name: Execute Jupyter Notebook
+        env:
+          ENV: 'dev'
+          LLM_API_KEY: ${{ secrets.OPENAI_API_KEY }}
+          OPENAI_API_KEY: ${{ secrets.OPENAI_API_KEY }}
+          EMBEDDING_MODEL: ${{ secrets.EMBEDDING_MODEL }}
+          EMBEDDING_ENDPOINT: ${{ secrets.EMBEDDING_ENDPOINT }}
+          EMBEDDING_API_KEY: ${{ secrets.EMBEDDING_API_KEY }}
+          EMBEDDING_API_VERSION: ${{ secrets.EMBEDDING_API_VERSION }}
+          GRAPHISTRY_USERNAME: ${{ secrets.GRAPHISTRY_USERNAME }}
+          GRAPHISTRY_PASSWORD: ${{ secrets.GRAPHISTRY_PASSWORD }}
+        run: |
+          poetry run jupyter nbconvert \
+          --to notebook \
+          --execute notebooks/llama_index_cognee_integration.ipynb \
+          --output executed_notebook.ipynb \
+          --ExecutePreprocessor.timeout=1200

.github/workflows/test_multimedia_example.yaml

@@ -16,7 +16,7 @@ jobs:
       with:
         example-location: ./examples/python/multimedia_example.py
       secrets:
-        LLM_API_KEY: ${{ secrets.OPENAI_API_KEY }} # Use OpenAI until we deploy models to handle multimedia
+        LLM_API_KEY: ${{ secrets.OPENAI_API_KEY }}
         OPENAI_API_KEY: ${{ secrets.OPENAI_API_KEY }}
         GRAPHISTRY_USERNAME: ${{ secrets.GRAPHISTRY_USERNAME }}
         GRAPHISTRY_PASSWORD: ${{ secrets.GRAPHISTRY_PASSWORD }}

.github/workflows/test_simple_example.yml

@@ -16,13 +16,7 @@ jobs:
       with:
         example-location: ./examples/python/simple_example.py
       secrets:
-        LLM_MODEL: ${{ secrets.LLM_MODEL }}
+        LLM_API_KEY: ${{ secrets.OPENAI_API_KEY }}
-        LLM_ENDPOINT: ${{ secrets.LLM_ENDPOINT }}
+        OPENAI_API_KEY: ${{ secrets.OPENAI_API_KEY }}
-        LLM_API_KEY: ${{ secrets.LLM_API_KEY }}
-        LLM_API_VERSION: ${{ secrets.LLM_API_VERSION }}
-        EMBEDDING_MODEL: ${{ secrets.EMBEDDING_MODEL }}
-        EMBEDDING_ENDPOINT: ${{ secrets.EMBEDDING_ENDPOINT }}
-        EMBEDDING_API_KEY: ${{ secrets.EMBEDDING_API_KEY }}
-        EMBEDDING_API_VERSION: ${{ secrets.EMBEDDING_API_VERSION }}
         GRAPHISTRY_USERNAME: ${{ secrets.GRAPHISTRY_USERNAME }}
         GRAPHISTRY_PASSWORD: ${{ secrets.GRAPHISTRY_PASSWORD }}

README.md

@@ -170,7 +170,7 @@ async def main():
     print(f"Searching cognee for insights with query: '{query_text}'")
     # Query cognee for insights on the added text
     search_results = await cognee.search(
-        SearchType.INSIGHTS, query_text=query_text
+        query_text=query_text, query_type=SearchType.INSIGHTS
     )
 
     print("Search results:")

cognee/api/v1/search/search_v2.py

@@ -8,8 +8,8 @@ from cognee.modules.search.methods import search as search_function
 
 
 async def search(
-    query_type: SearchType,
     query_text: str,
+    query_type: SearchType = SearchType.GRAPH_COMPLETION,
     user: User = None,
     datasets: Union[list[str], str, None] = None,
 ) -> list:

cognee/tests/test_falkordb.py

@@ -50,19 +50,23 @@ async def main():
     random_node = (await vector_engine.search("entity.name", "AI"))[0]
     random_node_name = random_node.payload["text"]
 
-    search_results = await cognee.search(SearchType.INSIGHTS, query_text=random_node_name)
+    search_results = await cognee.search(
+        query_type=SearchType.INSIGHTS, query_text=random_node_name
+    )
     assert len(search_results) != 0, "The search results list is empty."
     print("\n\nExtracted sentences are:\n")
     for result in search_results:
         print(f"{result}\n")
 
-    search_results = await cognee.search(SearchType.CHUNKS, query_text=random_node_name)
+    search_results = await cognee.search(query_type=SearchType.CHUNKS, query_text=random_node_name)
     assert len(search_results) != 0, "The search results list is empty."
     print("\n\nExtracted chunks are:\n")
     for result in search_results:
         print(f"{result}\n")
 
-    search_results = await cognee.search(SearchType.SUMMARIES, query_text=random_node_name)
+    search_results = await cognee.search(
+        query_type=SearchType.SUMMARIES, query_text=random_node_name
+    )
     assert len(search_results) != 0, "Query related summaries don't exist."
     print("\nExtracted summaries are:\n")
     for result in search_results:

cognee/tests/test_library.py

@@ -47,19 +47,23 @@ async def main():
     random_node = (await vector_engine.search("entity_name", "AI"))[0]
     random_node_name = random_node.payload["text"]
 
-    search_results = await cognee.search(SearchType.INSIGHTS, query_text=random_node_name)
+    search_results = await cognee.search(
+        query_type=SearchType.INSIGHTS, query_text=random_node_name
+    )
     assert len(search_results) != 0, "The search results list is empty."
     print("\n\nExtracted sentences are:\n")
     for result in search_results:
         print(f"{result}\n")
 
-    search_results = await cognee.search(SearchType.CHUNKS, query_text=random_node_name)
+    search_results = await cognee.search(query_type=SearchType.CHUNKS, query_text=random_node_name)
     assert len(search_results) != 0, "The search results list is empty."
     print("\n\nExtracted chunks are:\n")
     for result in search_results:
         print(f"{result}\n")
 
-    search_results = await cognee.search(SearchType.SUMMARIES, query_text=random_node_name)
+    search_results = await cognee.search(
+        query_type=SearchType.SUMMARIES, query_text=random_node_name
+    )
     assert len(search_results) != 0, "Query related summaries don't exist."
     print("\nExtracted summaries are:\n")
     for result in search_results:

cognee/tests/test_milvus.py

@@ -58,19 +58,23 @@ async def main():
     random_node = (await vector_engine.search("entity_name", "Quantum computer"))[0]
     random_node_name = random_node.payload["text"]
 
-    search_results = await cognee.search(SearchType.INSIGHTS, query_text=random_node_name)
+    search_results = await cognee.search(
+        query_type=SearchType.INSIGHTS, query_text=random_node_name
+    )
     assert len(search_results) != 0, "The search results list is empty."
     print("\n\nExtracted INSIGHTS are:\n")
     for result in search_results:
         print(f"{result}\n")
 
-    search_results = await cognee.search(SearchType.CHUNKS, query_text=random_node_name)
+    search_results = await cognee.search(query_type=SearchType.CHUNKS, query_text=random_node_name)
     assert len(search_results) != 0, "The search results list is empty."
     print("\n\nExtracted CHUNKS are:\n")
     for result in search_results:
         print(f"{result}\n")
 
-    search_results = await cognee.search(SearchType.SUMMARIES, query_text=random_node_name)
+    search_results = await cognee.search(
+        query_type=SearchType.SUMMARIES, query_text=random_node_name
+    )
     assert len(search_results) != 0, "The search results list is empty."
     print("\nExtracted SUMMARIES are:\n")
     for result in search_results:

cognee/tests/test_neo4j.py

@@ -51,19 +51,23 @@ async def main():
     random_node = (await vector_engine.search("entity_name", "Quantum computer"))[0]
     random_node_name = random_node.payload["text"]
 
-    search_results = await cognee.search(SearchType.INSIGHTS, query_text=random_node_name)
+    search_results = await cognee.search(
+        query_type=SearchType.INSIGHTS, query_text=random_node_name
+    )
     assert len(search_results) != 0, "The search results list is empty."
     print("\n\nExtracted sentences are:\n")
     for result in search_results:
         print(f"{result}\n")
 
-    search_results = await cognee.search(SearchType.CHUNKS, query_text=random_node_name)
+    search_results = await cognee.search(query_type=SearchType.CHUNKS, query_text=random_node_name)
     assert len(search_results) != 0, "The search results list is empty."
     print("\n\nExtracted chunks are:\n")
     for result in search_results:
         print(f"{result}\n")
 
-    search_results = await cognee.search(SearchType.SUMMARIES, query_text=random_node_name)
+    search_results = await cognee.search(
+        query_type=SearchType.SUMMARIES, query_text=random_node_name
+    )
     assert len(search_results) != 0, "Query related summaries don't exist."
     print("\nExtracted summaries are:\n")
     for result in search_results:

cognee/tests/test_pgvector.py

@@ -126,21 +126,25 @@ async def main():
     random_node = (await vector_engine.search("entity_name", "Quantum computer"))[0]
     random_node_name = random_node.payload["text"]
 
-    search_results = await cognee.search(SearchType.INSIGHTS, query_text=random_node_name)
+    search_results = await cognee.search(
+        query_type=SearchType.INSIGHTS, query_text=random_node_name
+    )
     assert len(search_results) != 0, "The search results list is empty."
     print("\n\nExtracted sentences are:\n")
     for result in search_results:
         print(f"{result}\n")
 
     search_results = await cognee.search(
-        SearchType.CHUNKS, query_text=random_node_name, datasets=[dataset_name_2]
+        query_type=SearchType.CHUNKS, query_text=random_node_name, datasets=[dataset_name_2]
     )
     assert len(search_results) != 0, "The search results list is empty."
     print("\n\nExtracted chunks are:\n")
     for result in search_results:
         print(f"{result}\n")
 
-    search_results = await cognee.search(SearchType.SUMMARIES, query_text=random_node_name)
+    search_results = await cognee.search(
+        query_type=SearchType.SUMMARIES, query_text=random_node_name
+    )
     assert len(search_results) != 0, "Query related summaries don't exist."
     print("\n\nExtracted summaries are:\n")
     for result in search_results:

cognee/tests/test_qdrant.py

@@ -51,19 +51,23 @@ async def main():
     random_node = (await vector_engine.search("entity_name", "Quantum computer"))[0]
     random_node_name = random_node.payload["text"]
 
-    search_results = await cognee.search(SearchType.INSIGHTS, query_text=random_node_name)
+    search_results = await cognee.search(
+        query_type=SearchType.INSIGHTS, query_text=random_node_name
+    )
     assert len(search_results) != 0, "The search results list is empty."
     print("\n\nExtracted sentences are:\n")
     for result in search_results:
         print(f"{result}\n")
 
-    search_results = await cognee.search(SearchType.CHUNKS, query_text=random_node_name)
+    search_results = await cognee.search(query_type=SearchType.CHUNKS, query_text=random_node_name)
     assert len(search_results) != 0, "The search results list is empty."
     print("\n\nExtracted chunks are:\n")
     for result in search_results:
         print(f"{result}\n")
 
-    search_results = await cognee.search(SearchType.SUMMARIES, query_text=random_node_name)
+    search_results = await cognee.search(
+        query_type=SearchType.SUMMARIES, query_text=random_node_name
+    )
     assert len(search_results) != 0, "Query related summaries don't exist."
     print("\nExtracted summaries are:\n")
     for result in search_results:

cognee/tests/test_weaviate.py

@@ -51,19 +51,23 @@ async def main():
     random_node = (await vector_engine.search("entity_name", "Quantum computer"))[0]
     random_node_name = random_node.payload["text"]
 
-    search_results = await cognee.search(SearchType.INSIGHTS, query_text=random_node_name)
+    search_results = await cognee.search(
+        query_text=random_node_name, query_type=SearchType.INSIGHTS
+    )
     assert len(search_results) != 0, "The search results list is empty."
     print("\n\nExtracted sentences are:\n")
     for result in search_results:
         print(f"{result}\n")
 
-    search_results = await cognee.search(SearchType.CHUNKS, query_text=random_node_name)
+    search_results = await cognee.search(query_type=SearchType.CHUNKS, query_text=random_node_name)
     assert len(search_results) != 0, "The search results list is empty."
     print("\n\nExtracted chunks are:\n")
     for result in search_results:
         print(f"{result}\n")
 
-    search_results = await cognee.search(SearchType.SUMMARIES, query_text=random_node_name)
+    search_results = await cognee.search(
+        query_type=SearchType.SUMMARIES, query_text=random_node_name
+    )
     assert len(search_results) != 0, "Query related summaries don't exist."
     print("\nExtracted summaries are:\n")
     for result in search_results:

evals/qa_context_provider_utils.py

@@ -76,7 +76,9 @@ async def get_context_with_cognee(
 
     search_results = []
     for search_type in search_types:
-        raw_search_results = await cognee.search(search_type, query_text=instance["question"])
+        raw_search_results = await cognee.search(
+            query_type=search_type, query_text=instance["question"]
+        )
 
         if search_type == SearchType.INSIGHTS:
             res_list = [_insight_to_string(edge) for edge in raw_search_results]

evals/simple_rag_vs_cognee_eval.py

@@ -94,7 +94,7 @@ async def cognify_search_base_rag(content: str, context: str):
 async def cognify_search_graph(content: str, context: str):
     from cognee.api.v1.search import search, SearchType
 
-    results = await search(SearchType.INSIGHTS, query_text="Donald Trump")
+    results = await search(query_type=SearchType.INSIGHTS, query_text="Donald Trump")
     print("results", results)
     return results
 

examples/python/dynamic_steps_example.py

@@ -186,7 +186,7 @@ async def main(enable_steps):
     # Step 4: Query insights
     if enable_steps.get("retriever"):
         search_results = await cognee.search(
-            SearchType.GRAPH_COMPLETION, query_text="Who has experience in design tools?"
+            query_type=SearchType.GRAPH_COMPLETION, query_text="Who has experience in design tools?"
         )
         print(search_results)
 

examples/python/multimedia_example.py

@@ -37,7 +37,7 @@ async def main():
 
     # Query cognee for summaries of the data in the multimedia files
     search_results = await cognee.search(
-        SearchType.SUMMARIES,
+        query_type=SearchType.SUMMARIES,
         query_text="What is in the multimedia files?",
     )
 

examples/python/simple_example.py

@@ -51,7 +51,7 @@ async def main():
     query_text = "Tell me about NLP"
     print(f"Searching cognee for insights with query: '{query_text}'")
     # Query cognee for insights on the added text
-    search_results = await cognee.search(SearchType.INSIGHTS, query_text=query_text)
+    search_results = await cognee.search(query_type=SearchType.INSIGHTS, query_text=query_text)
 
     print("Search results:")
     # Display results

modal_deployment.py

@@ -27,7 +27,7 @@ async def entry(text: str, query: str):
     await cognee.prune.prune_system(metadata=True)
     await cognee.add(text)
     await cognee.cognify()
-    search_results = await cognee.search(SearchType.GRAPH_COMPLETION, query_text=query)
+    search_results = await cognee.search(query_type=SearchType.GRAPH_COMPLETION, query_text=query)
 
     return {
         "text": text,

notebooks/cognee_demo.ipynb

@@ -830,7 +830,7 @@
     "node = (await vector_engine.search(\"entity_name\", \"sarah.nguyen@example.com\"))[0]\n",
     "node_name = node.payload[\"text\"]\n",
     "\n",
-    "search_results = await cognee.search(SearchType.SUMMARIES, query_text = node_name)\n",
+    "search_results = await cognee.search(query_type=SearchType.SUMMARIES, query_text = node_name)\n",
     "print(\"\\n\\Extracted summaries are:\\n\")\n",
     "for result in search_results:\n",
     "    print(f\"{result}\\n\")"
@@ -851,7 +851,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "search_results = await cognee.search(SearchType.CHUNKS, query_text = node_name)\n",
+    "search_results = await cognee.search(query_type=SearchType.CHUNKS, query_text = node_name)\n",
     "print(\"\\n\\nExtracted chunks are:\\n\")\n",
     "for result in search_results:\n",
     "    print(f\"{result}\\n\")"
@@ -872,7 +872,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "search_results = await cognee.search(SearchType.INSIGHTS, query_text = node_name)\n",
+    "search_results = await cognee.search(query_type=SearchType.INSIGHTS, query_text = node_name)\n",
     "print(\"\\n\\nExtracted sentences are:\\n\")\n",
     "for result in search_results:\n",
     "    print(f\"{result}\\n\")"

notebooks/cognee_llama_index.ipynb

@@ -179,7 +179,7 @@
     "\n",
     "# Query cognee for summaries\n",
     "search_results = await cognee.search(\n",
-    "    SearchType.SUMMARIES, query_text=\"What are the main news discussed in the document?\"\n",
+    "    query_type=SearchType.SUMMARIES, query_text=\"What are the main news discussed in the document?\"\n",
     ")\n",
     "# Display search results\n",
     "print(\"\\n Summary of main news discussed:\\n\")\n",

notebooks/cognee_multimedia_demo.ipynb

@@ -137,7 +137,7 @@
     "\n",
     "# Query cognee for summaries of the data in the multimedia files\n",
     "search_results = await cognee.search(\n",
-    "    SearchType.SUMMARIES,\n",
+    "    query_type=SearchType.SUMMARIES,\n",
     "    query_text=\"What is in the multimedia files?\",\n",
     ")\n",
     "\n",

notebooks/hr_demo.ipynb

@@ -566,7 +566,7 @@
     "node = (await vector_engine.search(\"entity_name\", \"sarah.nguyen@example.com\"))[0]\n",
     "node_name = node.payload[\"text\"]\n",
     "\n",
-    "search_results = await cognee.search(SearchType.SUMMARIES, query_text=node_name)\n",
+    "search_results = await cognee.search(query_type=SearchType.SUMMARIES, query_text=node_name)\n",
     "print(\"\\n\\Extracted summaries are:\\n\")\n",
     "for result in search_results:\n",
     "    print(f\"{result}\\n\")"
@@ -587,7 +587,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "search_results = await cognee.search(SearchType.CHUNKS, query_text=node_name)\n",
+    "search_results = await cognee.search(query_type=SearchType.CHUNKS, query_text=node_name)\n",
     "print(\"\\n\\nExtracted chunks are:\\n\")\n",
     "for result in search_results:\n",
     "    print(f\"{result}\\n\")"
@@ -608,7 +608,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "search_results = await cognee.search(SearchType.INSIGHTS, query_text=node_name)\n",
+    "search_results = await cognee.search(query_type=SearchType.INSIGHTS, query_text=node_name)\n",
     "print(\"\\n\\nExtracted sentences are:\\n\")\n",
     "for result in search_results:\n",
     "    print(f\"{result}\\n\")"