From 9c3e2422f305a7da481b09b1bdd6756caa99d168 Mon Sep 17 00:00:00 2001
From: Igor Ilic <igorilic03@gmail.com>
Date: Fri, 13 Dec 2024 15:18:33 +0100
Subject: [PATCH] feat: Add compute search to cognee

Add compute search to cognee which makes searches human readable

Feature COG-656
---
 cognee/api/v1/search/search_v2.py             |  3 ++
 cognee/tasks/compute/__init__.py              |  1 +
 cognee/tasks/compute/exceptions/__init__.py   |  9 +++++
 cognee/tasks/compute/exceptions/exceptions.py | 11 ++++++
 cognee/tasks/compute/query_compute.py         | 36 +++++++++++++++++++
 5 files changed, 60 insertions(+)
 create mode 100644 cognee/tasks/compute/__init__.py
 create mode 100644 cognee/tasks/compute/exceptions/__init__.py
 create mode 100644 cognee/tasks/compute/exceptions/exceptions.py
 create mode 100644 cognee/tasks/compute/query_compute.py

diff --git a/cognee/api/v1/search/search_v2.py b/cognee/api/v1/search/search_v2.py
index d77aa5fa8..516c7b11d 100644
--- a/cognee/api/v1/search/search_v2.py
+++ b/cognee/api/v1/search/search_v2.py
@@ -14,11 +14,13 @@ from cognee.modules.users.permissions.methods import get_document_ids_for_user
 from cognee.tasks.chunks import query_chunks
 from cognee.tasks.graph import query_graph_connections
 from cognee.tasks.summarization import query_summaries
+from cognee.tasks.compute import query_compute
 
 class SearchType(Enum):
     SUMMARIES = "SUMMARIES"
     INSIGHTS = "INSIGHTS"
     CHUNKS = "CHUNKS"
+    COMPUTE = "COMPUTE"
 
 async def search(query_type: SearchType, query_text: str, user: User = None) -> list:
     if user is None:
@@ -50,6 +52,7 @@ async def specific_search(query_type: SearchType, query: str, user) -> list:
         SearchType.SUMMARIES: query_summaries,
         SearchType.INSIGHTS: query_graph_connections,
         SearchType.CHUNKS: query_chunks,
+        SearchType.COMPUTE: query_compute,
     }
 
     search_task = search_tasks.get(query_type)
diff --git a/cognee/tasks/compute/__init__.py b/cognee/tasks/compute/__init__.py
new file mode 100644
index 000000000..7f6ed58fd
--- /dev/null
+++ b/cognee/tasks/compute/__init__.py
@@ -0,0 +1 @@
+from .query_compute import query_compute
\ No newline at end of file
diff --git a/cognee/tasks/compute/exceptions/__init__.py b/cognee/tasks/compute/exceptions/__init__.py
new file mode 100644
index 000000000..5f80e6ecc
--- /dev/null
+++ b/cognee/tasks/compute/exceptions/__init__.py
@@ -0,0 +1,9 @@
+"""
+Custom exceptions for the Cognee API.
+
+This module defines a set of exceptions for handling various compute errors
+"""
+
+from .exceptions import (
+    NoRelevantDataFound,
+)
\ No newline at end of file
diff --git a/cognee/tasks/compute/exceptions/exceptions.py b/cognee/tasks/compute/exceptions/exceptions.py
new file mode 100644
index 000000000..9b64c01a6
--- /dev/null
+++ b/cognee/tasks/compute/exceptions/exceptions.py
@@ -0,0 +1,11 @@
+from cognee.exceptions import CogneeApiError
+from fastapi import status
+
+class NoRelevantDataFound(CogneeApiError):
+    def __init__(
+            self,
+            message: str = "Search did not find any data.",
+            name: str = "NoRelevantDataFound",
+            status_code=status.HTTP_404_NOT_FOUND,
+    ):
+        super().__init__(message, name, status_code)
\ No newline at end of file
diff --git a/cognee/tasks/compute/query_compute.py b/cognee/tasks/compute/query_compute.py
new file mode 100644
index 000000000..1686701fe
--- /dev/null
+++ b/cognee/tasks/compute/query_compute.py
@@ -0,0 +1,36 @@
+from cognee.infrastructure.databases.vector import get_vector_engine
+from cognee.tasks.compute.exceptions import NoRelevantDataFound
+from cognee.infrastructure.llm.get_llm_client import get_llm_client
+from cognee.infrastructure.llm.prompts import read_query_prompt, render_prompt
+
+
+async def query_compute(query: str) -> list:
+    """
+    Parameters:
+    - query (str): The query string to compute.
+
+    Returns:
+    - list: Answer to the query.
+    """
+    vector_engine = get_vector_engine()
+
+    found_chunks = await vector_engine.search("document_chunk_text", query, limit = 1)
+
+    if len(found_chunks) == 0:
+        raise NoRelevantDataFound
+
+    args = {
+        "question": query,
+        "context": found_chunks[0].payload["text"],
+    }
+    user_prompt = render_prompt("context_for_question.txt", args)
+    system_prompt = read_query_prompt("answer_hotpot_using_cognee_search.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,
+    )
+
+    return [computed_answer]