cognee 0.3.6__py3-none-any.whl → 0.3.7.dev1__py3-none-any.whl

cognee/modules/retrieval/graph_completion_cot_retriever.py

@@ -1,15 +1,41 @@
+import asyncio
+import json
 from typing import Optional, List, Type, Any
+from pydantic import BaseModel
 from cognee.modules.graph.cognee_graph.CogneeGraphElements import Edge
 from cognee.shared.logging_utils import get_logger
 
 from cognee.modules.retrieval.graph_completion_retriever import GraphCompletionRetriever
-from cognee.modules.retrieval.utils.completion import generate_completion
+from cognee.modules.retrieval.utils.completion import (
+    generate_structured_completion,
+    summarize_text,
+)
+from cognee.modules.retrieval.utils.session_cache import (
+    save_conversation_history,
+    get_conversation_history,
+)
 from cognee.infrastructure.llm.LLMGateway import LLMGateway
 from cognee.infrastructure.llm.prompts import render_prompt, read_query_prompt
+from cognee.context_global_variables import session_user
+from cognee.infrastructure.databases.cache.config import CacheConfig
 
 logger = get_logger()
 
 
+def _as_answer_text(completion: Any) -> str:
+    """Convert completion to human-readable text for validation and follow-up prompts."""
+    if isinstance(completion, str):
+        return completion
+    if isinstance(completion, BaseModel):
+        # Add notice that this is a structured response
+        json_str = completion.model_dump_json(indent=2)
+        return f"[Structured Response]\n{json_str}"
+    try:
+        return json.dumps(completion, indent=2)
+    except TypeError:
+        return str(completion)
+
+
 class GraphCompletionCotRetriever(GraphCompletionRetriever):
     """
     Handles graph completion by generating responses based on a series of interactions with
@@ -18,6 +44,7 @@ class GraphCompletionCotRetriever(GraphCompletionRetriever):
     questions based on reasoning. The public methods are:
 
     - get_completion
+    - get_structured_completion
 
     Instance variables include:
     - validation_system_prompt_path
@@ -54,33 +81,30 @@ class GraphCompletionCotRetriever(GraphCompletionRetriever):
         self.followup_system_prompt_path = followup_system_prompt_path
         self.followup_user_prompt_path = followup_user_prompt_path
 
-    async def get_completion(
+    async def _run_cot_completion(
         self,
         query: str,
         context: Optional[List[Edge]] = None,
-        max_iter=4,
-    ) -> List[str]:
+        conversation_history: str = "",
+        max_iter: int = 4,
+        response_model: Type = str,
+    ) -> tuple[Any, str, List[Edge]]:
         """
-        Generate completion responses based on a user query and contextual information.
-
-        This method interacts with a language model client to retrieve a structured response,
-        using a series of iterations to refine the answers and generate follow-up questions
-        based on reasoning derived from previous outputs. It raises exceptions if the context
-        retrieval fails or if the model encounters issues in generating outputs.
+        Run chain-of-thought completion with optional structured output.
 
         Parameters:
         -----------
-
-            - query (str): The user's query to be processed and answered.
-            - context (Optional[Any]): Optional context that may assist in answering the query.
-              If not provided, it will be fetched based on the query. (default None)
-            - max_iter: The maximum number of iterations to refine the answer and generate
-              follow-up questions. (default 4)
+            - query: User query
+            - context: Optional pre-fetched context edges
+            - conversation_history: Optional conversation history string
+            - max_iter: Maximum CoT iterations
+            - response_model: Type for structured output (str for plain text)
 
         Returns:
         --------
-
-            - List[str]: A list containing the generated answer to the user's query.
+            - completion_result: The generated completion (string or structured model)
+            - context_text: The resolved context text
+            - triplets: The list of triplets used
         """
         followup_question = ""
         triplets = []
@@ -97,16 +121,21 @@ class GraphCompletionCotRetriever(GraphCompletionRetriever):
                 triplets += await self.get_context(followup_question)
                 context_text = await self.resolve_edges_to_text(list(set(triplets)))
 
-            completion = await generate_completion(
+            completion = await generate_structured_completion(
                 query=query,
                 context=context_text,
                 user_prompt_path=self.user_prompt_path,
                 system_prompt_path=self.system_prompt_path,
                 system_prompt=self.system_prompt,
+                conversation_history=conversation_history if conversation_history else None,
+                response_model=response_model,
             )
+
             logger.info(f"Chain-of-thought: round {round_idx} - answer: {completion}")
+
             if round_idx < max_iter:
-                valid_args = {"query": query, "answer": completion, "context": context_text}
+                answer_text = _as_answer_text(completion)
+                valid_args = {"query": query, "answer": answer_text, "context": context_text}
                 valid_user_prompt = render_prompt(
                     filename=self.validation_user_prompt_path, context=valid_args
                 )
@@ -119,7 +148,7 @@ class GraphCompletionCotRetriever(GraphCompletionRetriever):
                     system_prompt=valid_system_prompt,
                     response_model=str,
                 )
-                followup_args = {"query": query, "answer": completion, "reasoning": reasoning}
+                followup_args = {"query": query, "answer": answer_text, "reasoning": reasoning}
                 followup_prompt = render_prompt(
                     filename=self.followup_user_prompt_path, context=followup_args
                 )
@@ -134,9 +163,110 @@ class GraphCompletionCotRetriever(GraphCompletionRetriever):
                     f"Chain-of-thought: round {round_idx} - follow-up question: {followup_question}"
                 )
 
+        return completion, context_text, triplets
+
+    async def get_structured_completion(
+        self,
+        query: str,
+        context: Optional[List[Edge]] = None,
+        session_id: Optional[str] = None,
+        max_iter: int = 4,
+        response_model: Type = str,
+    ) -> Any:
+        """
+        Generate structured completion responses based on a user query and contextual information.
+
+        This method applies the same chain-of-thought logic as get_completion but returns
+        structured output using the provided response model.
+
+        Parameters:
+        -----------
+            - query (str): The user's query to be processed and answered.
+            - context (Optional[List[Edge]]): Optional context that may assist in answering the query.
+              If not provided, it will be fetched based on the query. (default None)
+            - session_id (Optional[str]): Optional session identifier for caching. If None,
+              defaults to 'default_session'. (default None)
+            - max_iter: The maximum number of iterations to refine the answer and generate
+              follow-up questions. (default 4)
+            - response_model (Type): The Pydantic model type for structured output. (default str)
+
+        Returns:
+        --------
+            - Any: The generated structured completion based on the response model.
+        """
+        # Check if session saving is enabled
+        cache_config = CacheConfig()
+        user = session_user.get()
+        user_id = getattr(user, "id", None)
+        session_save = user_id and cache_config.caching
+
+        # Load conversation history if enabled
+        conversation_history = ""
+        if session_save:
+            conversation_history = await get_conversation_history(session_id=session_id)
+
+        completion, context_text, triplets = await self._run_cot_completion(
+            query=query,
+            context=context,
+            conversation_history=conversation_history,
+            max_iter=max_iter,
+            response_model=response_model,
+        )
+
         if self.save_interaction and context and triplets and completion:
             await self.save_qa(
-                question=query, answer=completion, context=context_text, triplets=triplets
+                question=query, answer=str(completion), context=context_text, triplets=triplets
+            )
+
+        # Save to session cache if enabled
+        if session_save:
+            context_summary = await summarize_text(context_text)
+            await save_conversation_history(
+                query=query,
+                context_summary=context_summary,
+                answer=str(completion),
+                session_id=session_id,
             )
 
+        return completion
+
+    async def get_completion(
+        self,
+        query: str,
+        context: Optional[List[Edge]] = None,
+        session_id: Optional[str] = None,
+        max_iter=4,
+    ) -> List[str]:
+        """
+        Generate completion responses based on a user query and contextual information.
+
+        This method interacts with a language model client to retrieve a structured response,
+        using a series of iterations to refine the answers and generate follow-up questions
+        based on reasoning derived from previous outputs. It raises exceptions if the context
+        retrieval fails or if the model encounters issues in generating outputs.
+
+        Parameters:
+        -----------
+
+            - query (str): The user's query to be processed and answered.
+            - context (Optional[Any]): Optional context that may assist in answering the query.
+              If not provided, it will be fetched based on the query. (default None)
+            - session_id (Optional[str]): Optional session identifier for caching. If None,
+              defaults to 'default_session'. (default None)
+            - max_iter: The maximum number of iterations to refine the answer and generate
+              follow-up questions. (default 4)
+
+        Returns:
+        --------
+
+            - List[str]: A list containing the generated answer to the user's query.
+        """
+        completion = await self.get_structured_completion(
+            query=query,
+            context=context,
+            session_id=session_id,
+            max_iter=max_iter,
+            response_model=str,
+        )
+
         return [completion]

cognee/modules/retrieval/graph_completion_retriever.py

@@ -1,20 +1,26 @@
+import asyncio
 from typing import Any, Optional, Type, List
 from uuid import NAMESPACE_OID, uuid5
 
 from cognee.infrastructure.engine import DataPoint
 from cognee.modules.graph.cognee_graph.CogneeGraphElements import Edge
-from cognee.modules.users.methods import get_default_user
 from cognee.tasks.storage import add_data_points
 from cognee.modules.graph.utils import resolve_edges_to_text
 from cognee.modules.graph.utils.convert_node_to_data_point import get_all_subclasses
 from cognee.modules.retrieval.base_graph_retriever import BaseGraphRetriever
 from cognee.modules.retrieval.utils.brute_force_triplet_search import brute_force_triplet_search
-from cognee.modules.retrieval.utils.completion import generate_completion
+from cognee.modules.retrieval.utils.completion import generate_completion, summarize_text
+from cognee.modules.retrieval.utils.session_cache import (
+    save_conversation_history,
+    get_conversation_history,
+)
 from cognee.shared.logging_utils import get_logger
 from cognee.modules.retrieval.utils.extract_uuid_from_node import extract_uuid_from_node
 from cognee.modules.retrieval.utils.models import CogneeUserInteraction
 from cognee.modules.engine.models.node_set import NodeSet
 from cognee.infrastructure.databases.graph import get_graph_engine
+from cognee.context_global_variables import session_user
+from cognee.infrastructure.databases.cache.config import CacheConfig
 
 logger = get_logger("GraphCompletionRetriever")
 
@@ -118,6 +124,13 @@ class GraphCompletionRetriever(BaseGraphRetriever):
             - str: A string representing the resolved context from the retrieved triplets, or an
               empty string if no triplets are found.
         """
+        graph_engine = await get_graph_engine()
+        is_empty = await graph_engine.is_empty()
+
+        if is_empty:
+            logger.warning("Search attempt on an empty knowledge graph")
+            return []
+
         triplets = await self.get_triplets(query)
 
         if len(triplets) == 0:
@@ -132,6 +145,7 @@ class GraphCompletionRetriever(BaseGraphRetriever):
         self,
         query: str,
         context: Optional[List[Edge]] = None,
+        session_id: Optional[str] = None,
     ) -> List[str]:
         """
         Generates a completion using graph connections context based on a query.
@@ -142,6 +156,8 @@ class GraphCompletionRetriever(BaseGraphRetriever):
             - query (str): The query string for which a completion is generated.
             - context (Optional[Any]): Optional context to use for generating the completion; if
               not provided, context is retrieved based on the query. (default None)
+            - session_id (Optional[str]): Optional session identifier for caching. If None,
+              defaults to 'default_session'. (default None)
 
         Returns:
         --------
@@ -155,19 +171,47 @@ class GraphCompletionRetriever(BaseGraphRetriever):
 
         context_text = await resolve_edges_to_text(triplets)
 
-        completion = await generate_completion(
-            query=query,
-            context=context_text,
-            user_prompt_path=self.user_prompt_path,
-            system_prompt_path=self.system_prompt_path,
-            system_prompt=self.system_prompt,
-        )
+        cache_config = CacheConfig()
+        user = session_user.get()
+        user_id = getattr(user, "id", None)
+        session_save = user_id and cache_config.caching
+
+        if session_save:
+            conversation_history = await get_conversation_history(session_id=session_id)
+
+            context_summary, completion = await asyncio.gather(
+                summarize_text(context_text),
+                generate_completion(
+                    query=query,
+                    context=context_text,
+                    user_prompt_path=self.user_prompt_path,
+                    system_prompt_path=self.system_prompt_path,
+                    system_prompt=self.system_prompt,
+                    conversation_history=conversation_history,
+                ),
+            )
+        else:
+            completion = await generate_completion(
+                query=query,
+                context=context_text,
+                user_prompt_path=self.user_prompt_path,
+                system_prompt_path=self.system_prompt_path,
+                system_prompt=self.system_prompt,
+            )
 
         if self.save_interaction and context and triplets and completion:
             await self.save_qa(
                 question=query, answer=completion, context=context_text, triplets=triplets
             )
 
+        if session_save:
+            await save_conversation_history(
+                query=query,
+                context_summary=context_summary,
+                answer=completion,
+                session_id=session_id,
+            )
+
         return [completion]
 
     async def save_qa(self, question: str, answer: str, context: str, triplets: List) -> None:
@@ -194,7 +238,7 @@ class GraphCompletionRetriever(BaseGraphRetriever):
             belongs_to_set=interactions_node_set,
         )
 
-        await add_data_points(data_points=[cognee_user_interaction], update_edge_collection=False)
+        await add_data_points(data_points=[cognee_user_interaction])
 
         relationships = []
         relationship_name = "used_graph_element_to_answer"

cognee/modules/retrieval/lexical_retriever.py

@@ -116,8 +116,26 @@ class LexicalRetriever(BaseRetriever):
         else:
             return [self.payloads[chunk_id] for chunk_id, _ in top_results]
 
-    async def get_completion(self, query: str, context: Optional[Any] = None) -> Any:
-        """Returns context for the given query (retrieves if not provided)."""
+    async def get_completion(
+        self, query: str, context: Optional[Any] = None, session_id: Optional[str] = None
+    ) -> Any:
+        """
+        Returns context for the given query (retrieves if not provided).
+
+        Parameters:
+        -----------
+
+            - query (str): The query string to retrieve context for.
+            - context (Optional[Any]): Optional pre-fetched context; if None, it retrieves
+              the context for the query. (default None)
+            - session_id (Optional[str]): Optional session identifier for caching. If None,
+              defaults to 'default_session'. (default None)
+
+        Returns:
+        --------
+
+            - Any: The context, either provided or retrieved.
+        """
         if context is None:
             context = await self.get_context(query)
         return context

cognee/modules/retrieval/natural_language_retriever.py

@@ -122,10 +122,17 @@ class NaturalLanguageRetriever(BaseRetriever):
               query.
         """
         graph_engine = await get_graph_engine()
+        is_empty = await graph_engine.is_empty()
+
+        if is_empty:
+            logger.warning("Search attempt on an empty knowledge graph")
+            return []
 
         return await self._execute_cypher_query(query, graph_engine)
 
-    async def get_completion(self, query: str, context: Optional[Any] = None) -> Any:
+    async def get_completion(
+        self, query: str, context: Optional[Any] = None, session_id: Optional[str] = None
+    ) -> Any:
         """
         Returns a completion based on the query and context.
 
@@ -139,6 +146,8 @@ class NaturalLanguageRetriever(BaseRetriever):
             - query (str): The natural language query to get a completion from.
             - context (Optional[Any]): The context in which to base the completion; if not
               provided, it will be retrieved using the query. (default None)
+            - session_id (Optional[str]): Optional session identifier for caching. If None,
+              defaults to 'default_session'. (default None)
 
         Returns:
         --------

cognee/modules/retrieval/summaries_retriever.py

@@ -62,7 +62,9 @@ class SummariesRetriever(BaseRetriever):
         logger.info(f"Returning {len(summary_payloads)} summary payloads")
         return summary_payloads
 
-    async def get_completion(self, query: str, context: Optional[Any] = None, **kwargs) -> Any:
+    async def get_completion(
+        self, query: str, context: Optional[Any] = None, session_id: Optional[str] = None, **kwargs
+    ) -> Any:
         """
         Generates a completion using summaries context.
 
@@ -75,6 +77,8 @@ class SummariesRetriever(BaseRetriever):
             - query (str): The search query for generating the completion.
             - context (Optional[Any]): Optional context for the completion; if not provided,
               will be retrieved based on the query. (default None)
+            - session_id (Optional[str]): Optional session identifier for caching. If None,
+              defaults to 'default_session'. (default None)
 
         Returns:
         --------

cognee/modules/retrieval/temporal_retriever.py

@@ -1,16 +1,22 @@
 import os
+import asyncio
 from typing import Any, Optional, List, Type
 
 
 from operator import itemgetter
 from cognee.infrastructure.databases.vector import get_vector_engine
-from cognee.modules.retrieval.utils.completion import generate_completion
+from cognee.modules.retrieval.utils.completion import generate_completion, summarize_text
+from cognee.modules.retrieval.utils.session_cache import (
+    save_conversation_history,
+    get_conversation_history,
+)
 from cognee.infrastructure.databases.graph import get_graph_engine
 from cognee.infrastructure.llm.prompts import render_prompt
 from cognee.infrastructure.llm import LLMGateway
 from cognee.modules.retrieval.graph_completion_retriever import GraphCompletionRetriever
 from cognee.shared.logging_utils import get_logger
-
+from cognee.context_global_variables import session_user
+from cognee.infrastructure.databases.cache.config import CacheConfig
 
 from cognee.tasks.temporal_graph.models import QueryInterval
 
@@ -137,17 +143,63 @@ class TemporalRetriever(GraphCompletionRetriever):
 
         return self.descriptions_to_string(top_k_events)
 
-    async def get_completion(self, query: str, context: Optional[str] = None) -> List[str]:
-        """Generates a response using the query and optional context."""
+    async def get_completion(
+        self, query: str, context: Optional[str] = None, session_id: Optional[str] = None
+    ) -> List[str]:
+        """
+        Generates a response using the query and optional context.
+
+        Parameters:
+        -----------
+
+            - query (str): The query string for which a completion is generated.
+            - context (Optional[str]): Optional context to use; if None, it will be
+              retrieved based on the query. (default None)
+            - session_id (Optional[str]): Optional session identifier for caching. If None,
+              defaults to 'default_session'. (default None)
+
+        Returns:
+        --------
+
+            - List[str]: A list containing the generated completion.
+        """
         if not context:
             context = await self.get_context(query=query)
 
         if context:
-            completion = await generate_completion(
-                query=query,
-                context=context,
-                user_prompt_path=self.user_prompt_path,
-                system_prompt_path=self.system_prompt_path,
-            )
+            # Check if we need to generate context summary for caching
+            cache_config = CacheConfig()
+            user = session_user.get()
+            user_id = getattr(user, "id", None)
+            session_save = user_id and cache_config.caching
+
+            if session_save:
+                conversation_history = await get_conversation_history(session_id=session_id)
+
+                context_summary, completion = await asyncio.gather(
+                    summarize_text(context),
+                    generate_completion(
+                        query=query,
+                        context=context,
+                        user_prompt_path=self.user_prompt_path,
+                        system_prompt_path=self.system_prompt_path,
+                        conversation_history=conversation_history,
+                    ),
+                )
+            else:
+                completion = await generate_completion(
+                    query=query,
+                    context=context,
+                    user_prompt_path=self.user_prompt_path,
+                    system_prompt_path=self.system_prompt_path,
+                )
+
+            if session_save:
+                await save_conversation_history(
+                    query=query,
+                    context_summary=context_summary,
+                    answer=completion,
+                    session_id=session_id,
+                )
 
         return [completion]

cognee/modules/retrieval/user_qa_feedback.py

@@ -8,7 +8,7 @@ from cognee.shared.logging_utils import get_logger
 from cognee.modules.retrieval.base_feedback import BaseFeedback
 from cognee.modules.retrieval.utils.models import CogneeUserFeedback
 from cognee.modules.retrieval.utils.models import UserFeedbackEvaluation
-from cognee.tasks.storage import add_data_points
+from cognee.tasks.storage import add_data_points, index_graph_edges
 
 logger = get_logger("CompletionRetriever")
 
@@ -47,7 +47,7 @@ class UserQAFeedback(BaseFeedback):
             belongs_to_set=feedbacks_node_set,
         )
 
-        await add_data_points(data_points=[cognee_user_feedback], update_edge_collection=False)
+        await add_data_points(data_points=[cognee_user_feedback])
 
         relationships = []
         relationship_name = "gives_feedback_to"
@@ -76,6 +76,7 @@ class UserQAFeedback(BaseFeedback):
         if len(relationships) > 0:
             graph_engine = await get_graph_engine()
             await graph_engine.add_edges(relationships)
+            await index_graph_edges(relationships)
             await graph_engine.apply_feedback_weight(
                 node_ids=to_node_ids, weight=feedback_sentiment.score
             )

cognee/modules/retrieval/utils/completion.py

@@ -1,23 +1,49 @@
-from typing import Optional
+from typing import Optional, Type, Any
 from cognee.infrastructure.llm.LLMGateway import LLMGateway
 from cognee.infrastructure.llm.prompts import render_prompt, read_query_prompt
 
 
-async def generate_completion(
+async def generate_structured_completion(
     query: str,
     context: str,
     user_prompt_path: str,
     system_prompt_path: str,
     system_prompt: Optional[str] = None,
-) -> str:
-    """Generates a completion using LLM with given context and prompts."""
+    conversation_history: Optional[str] = None,
+    response_model: Type = str,
+) -> Any:
+    """Generates a structured completion using LLM with given context and prompts."""
     args = {"question": query, "context": context}
     user_prompt = render_prompt(user_prompt_path, args)
     system_prompt = system_prompt if system_prompt else read_query_prompt(system_prompt_path)
 
+    if conversation_history:
+        #:TODO: I would separate the history and put it into the system prompt but we have to test what works best with longer convos
+        system_prompt = conversation_history + "\nTASK:" + system_prompt
+
     return await LLMGateway.acreate_structured_output(
         text_input=user_prompt,
         system_prompt=system_prompt,
+        response_model=response_model,
+    )
+
+
+async def generate_completion(
+    query: str,
+    context: str,
+    user_prompt_path: str,
+    system_prompt_path: str,
+    system_prompt: Optional[str] = None,
+    conversation_history: Optional[str] = None,
+) -> str:
+    """Generates a completion using LLM with given context and prompts."""
+    return await generate_structured_completion(
+        query=query,
+        context=context,
+        user_prompt_path=user_prompt_path,
+        system_prompt_path=system_prompt_path,
+        system_prompt=system_prompt,
+        conversation_history=conversation_history,
         response_model=str,
     )
 

cognee/modules/retrieval/utils/description_to_codepart_search.py

@@ -62,7 +62,7 @@ async def code_description_to_code_part(
 
     try:
         if include_docs:
-            search_results = await search(query_text=query, query_type="INSIGHTS")
+            search_results = await search(query_text=query, query_type="GRAPH_COMPLETION")
 
             concatenated_descriptions = " ".join(
                 obj["description"]