cognee 1.0.1.dev2__tar.gz → 1.0.1.dev3__tar.gz

{cognee-1.0.1.dev2 → cognee-1.0.1.dev3}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: cognee
-Version: 1.0.1.dev2
+Version: 1.0.1.dev3
 Summary: Cognee - is a library for enriching LLM context with a semantic layer for better understanding and reasoning.
 Project-URL: Homepage, https://www.cognee.ai
 Project-URL: Repository, https://github.com/topoteretes/cognee

{cognee-1.0.1.dev2 → cognee-1.0.1.dev3}/cognee/api/v1/cognify/cognify.py

@@ -1,6 +1,6 @@
 import asyncio
 from pydantic import BaseModel
-from typing import Union, Optional, List
+from typing import Union, Optional
 from uuid import UUID
 
 from cognee.modules.cognify.config import get_cognify_config
@@ -23,9 +23,8 @@ from cognee.tasks.documents import (
     classify_documents,
     extract_chunks_from_documents,
 )
-from cognee.tasks.graph import extract_graph_from_data
+from cognee.tasks.graph.extract_graph_and_summarize import extract_graph_and_summarize
 from cognee.tasks.storage import add_data_points
-from cognee.tasks.summarization import summarize_text
 from cognee.tasks.ingestion.extract_dlt_fk_edges import extract_dlt_fk_edges
 from cognee.modules.pipelines.layers.pipeline_execution_mode import get_pipeline_executor
 from cognee.tasks.temporal_graph.extract_events_and_entities import extract_events_and_timestamps
@@ -318,19 +317,15 @@ async def get_default_tasks(  # TODO: Find out a better way to do this (Boris's
             chunker=chunker,
         ),
         # COGNIFY: LLM-extract entities and relationships into a knowledge graph
+        # COGNIFY: LLM-summarize each chunk for hierarchical retrieval
         Task(
-            extract_graph_from_data,
+            extract_graph_and_summarize,
             graph_model=graph_model,
             config=config,
             custom_prompt=custom_prompt,
             task_config={"batch_size": chunks_per_batch},
             **kwargs,
         ),
-        # COGNIFY: LLM-summarize each chunk for hierarchical retrieval
-        Task(
-            summarize_text,
-            task_config={"batch_size": chunks_per_batch},
-        ),
         # LOAD: persist nodes, edges, and embeddings to graph/vector DBs
         Task(
             add_data_points,

{cognee-1.0.1.dev2 → cognee-1.0.1.dev3}/cognee/api/v1/llm/routers/get_llm_router.py

@@ -1,7 +1,7 @@
 import json
-from typing import Any, Dict
+from typing import Any, Dict, List, Annotated
 
-from fastapi import APIRouter, Depends
+from fastapi import APIRouter, Depends, File, UploadFile as UF, Form
 from fastapi.responses import JSONResponse
 from pydantic import Field
 
@@ -9,14 +9,24 @@ from cognee import __version__ as cognee_version
 from cognee.api.DTO import InDTO, OutDTO
 from cognee.infrastructure.llm.LLMGateway import LLMGateway
 from cognee.infrastructure.llm.prompts import render_prompt
+from cognee.infrastructure.loaders import get_loader_engine
 from cognee.modules.users.methods import get_authenticated_user
 from cognee.modules.users.models import User
 from cognee.shared.logging_utils import get_logger
 from cognee.shared.usage_logger import log_usage
 from cognee.shared.utils import send_telemetry
+from pydantic import WithJsonSchema
+from contextlib import asynccontextmanager
+from pathlib import Path
+import os
+import tempfile
 
 logger = get_logger("api.llm")
 
+# NOTE: Needed because of: https://github.com/fastapi/fastapi/discussions/14975
+#       Once issue is resolved on Swagger side it can be removed.
+UploadFile = Annotated[UF, WithJsonSchema({"type": "string", "format": "binary"})]
+
 _ALLOWED_LLM_PARAMS = {"temperature", "max_tokens", "top_p", "seed"}
 
 
@@ -24,6 +34,31 @@ def _safe_params(params: dict) -> dict:
     return {k: v for k, v in params.items() if k in _ALLOWED_LLM_PARAMS}
 
 
+@asynccontextmanager
+async def upload_to_temp_path(upload: UploadFile):
+    """
+    Materialize UploadFile to a real temp file path for path-based loaders.
+    Cleans up after use.
+    """
+    suffix = Path(upload.filename or "upload.bin").suffix or ".bin"
+
+    fd, temp_path = tempfile.mkstemp(prefix="cognee_upload_", suffix=suffix)
+    os.close(fd)  # close low-level fd; we'll write normally
+
+    try:
+        raw = await upload.read()
+        with open(temp_path, "wb") as f:
+            f.write(raw)
+
+        await upload.seek(0)  # keep UploadFile reusable
+        yield temp_path
+    finally:
+        try:
+            os.remove(temp_path)
+        except FileNotFoundError:
+            pass
+
+
 class CustomPromptGenerationPayloadDTO(InDTO):
     graph_model: Dict[str, Any] = Field(..., description="Graph model schema as JSON object.")
     parameters: Dict[str, Any] = Field(
@@ -36,14 +71,6 @@ class CustomPromptGenerationResponseDTO(OutDTO):
     custom_prompt: str
 
 
-class InferSchemaPayloadDTO(InDTO):
-    text: str = Field(..., description="Sample text to analyze for entity types and relationships.")
-    parameters: Dict[str, Any] = Field(
-        default_factory=dict,
-        description="Additional kwargs forwarded to LLMGateway.",
-    )
-
-
 class InferSchemaResponseDTO(OutDTO):
     graph_schema: Dict[str, Any]
 
@@ -101,11 +128,16 @@ def get_llm_router() -> APIRouter:
     @router.post("/infer-schema", response_model=InferSchemaResponseDTO)
     @log_usage(function_name="POST /v1/llm/infer-schema", log_type="api_endpoint")
     async def infer_schema(
-        payload: InferSchemaPayloadDTO,
+        data: List[UploadFile] = File(default=None),
+        text: str = Form(default=None),
+        parameters: Dict[str, Any] = Form(
+            default_factory=dict,
+            description="Additional kwargs forwarded to LLMGateway.",
+        ),
         user: User = Depends(get_authenticated_user),
     ):
         """
-        Analyze sample text and propose a JSON Schema describing the entity types
+        Analyze sample text and/or uploaded files, and propose a JSON Schema describing the entity types
         and relationships present. The returned schema can be passed directly to
         ``/v1/llm/custom-prompt`` or ``/v1/cognify``.
         """
@@ -114,15 +146,42 @@ def get_llm_router() -> APIRouter:
             user.id,
             additional_properties={
                 "endpoint": "POST /v1/llm/infer-schema",
-                "text_length": len(payload.text),
+                "filenames": [f.filename for f in (data or [])],
+                "text_length": len(text or ""),
                 "cognee_version": cognee_version,
             },
         )
 
+        if not data and not text:
+            return JSONResponse(
+                status_code=400,
+                content={"error": "Either text or at least one file must be provided."},
+            )
         try:
+            file_contents = []
+            for file in data or []:
+                async with upload_to_temp_path(file) as file_path:
+                    loader = get_loader_engine().get_loader(
+                        file_path=file_path, preferred_loaders={}
+                    )
+                    if not loader:
+                        return JSONResponse(
+                            status_code=400,
+                            content={
+                                "error": f"No valid loader found for file {file.filename}. File type is not supported."
+                            },
+                        )
+                    data_content = await loader.load(file_path, persist=False)
+                    file_contents.append(data_content)
+
+            if text:
+                file_contents.append(text)
+
+            file_contents_text = "\n\n".join(file_contents)
+
             user_prompt = render_prompt(
                 "infer_schema_user.txt",
-                {"SAMPLE_TEXT": payload.text},
+                {"SAMPLE_TEXT": file_contents_text},
             )
 
             system_prompt = render_prompt(
@@ -134,7 +193,7 @@ def get_llm_router() -> APIRouter:
                 text_input=user_prompt,
                 system_prompt=system_prompt,
                 response_model=str,  # type: ignore[arg-type]
-                **_safe_params(payload.parameters),
+                **_safe_params(parameters),
             )
 
             # Parse the LLM output as JSON

{cognee-1.0.1.dev2 → cognee-1.0.1.dev3}/cognee/api/v1/remember/remember.py

@@ -12,6 +12,9 @@ from typing_extensions import TypedDict
 
 from cognee.shared.logging_utils import get_logger
 from cognee.tasks.ingestion.data_item import DataItem
+from cognee.modules.pipelines.layers.resolve_authorized_user_datasets import (
+    resolve_authorized_user_datasets,
+)
 from cognee.modules.observability import (
     new_span,
     COGNEE_DATASET_NAME,
@@ -542,6 +545,10 @@ async def _remember_inner(
         return result
 
     # Build the result object — starts as "running"
+    if not dataset_id and dataset_name:
+        # Create dataset if it doesn't exist
+        user, dataset_id = await resolve_authorized_user_datasets(dataset_name, user)
+        dataset_id = dataset_id[0].id if dataset_id else None
     result = RememberResult(
         status="running",
         dataset_name=dataset_name,

{cognee-1.0.1.dev2 → cognee-1.0.1.dev3}/cognee/api/v1/remember/routers/get_remember_router.py

@@ -30,7 +30,7 @@ def get_remember_router() -> APIRouter:
         data: List[UploadFile] = File(default=None),
         datasetName: Optional[str] = Form(default=None),
         datasetId: Union[UUID, Literal[""], None] = Form(default=None, examples=[""]),
-        session_id: Optional[str] = Form(default=None),
+        session_id: Optional[str] = Form(default=None, examples=[""]),
         node_set: Optional[List[str]] = Form(default=[""], example=[""]),
         run_in_background: Optional[bool] = Form(default=False),
         custom_prompt: Optional[str] = Form(default=""),

{cognee-1.0.1.dev2 → cognee-1.0.1.dev3}/cognee/api/v1/search/routers/get_search_router.py

@@ -155,7 +155,9 @@ def get_search_router() -> APIRouter:
                 query_text=payload.query,
                 query_type=payload.search_type,
                 user=user,
-                datasets=payload.datasets,
+                datasets=payload.datasets
+                if not payload.dataset_ids
+                else None,  # If dataset_ids are provided, ignore datasets by name to avoid confusion and potential mismatches.
                 dataset_ids=payload.dataset_ids,
                 system_prompt=payload.system_prompt,
                 node_name=payload.node_name,

{cognee-1.0.1.dev2 → cognee-1.0.1.dev3}/cognee/infrastructure/databases/cache/cache_db_interface.py

@@ -152,6 +152,40 @@ class CacheDBInterface(ABC):
         """
         pass
 
+    @abstractmethod
+    async def append_agent_trace_step(
+        self,
+        user_id: str,
+        session_id: str,
+        trace_id: str,
+        origin_function: str,
+        status: str,
+        memory_query: str = "",
+        memory_context: str = "",
+        method_params: dict | None = None,
+        method_return_value=None,
+        error_message: str = "",
+        session_feedback: str = "",
+    ) -> None:
+        """
+        Append one agent trace step to the session-scoped trace list.
+        """
+        pass
+
+    @abstractmethod
+    async def get_agent_trace_session(self, user_id: str, session_id: str) -> list[dict]:
+        """
+        Retrieve all agent trace steps for the given session.
+        """
+        pass
+
+    @abstractmethod
+    async def get_agent_trace_feedback(self, user_id: str, session_id: str) -> list[str]:
+        """
+        Retrieve only per-step feedback strings for the given trace session.
+        """
+        pass
+
     @abstractmethod
     async def prune(self) -> None:
         """

{cognee-1.0.1.dev2 → cognee-1.0.1.dev3}/cognee/infrastructure/databases/cache/fscache/FsCacheAdapter.py

@@ -7,7 +7,7 @@ import diskcache as dc
 from pydantic import ValidationError
 
 from cognee.infrastructure.databases.cache.cache_db_interface import CacheDBInterface
-from cognee.infrastructure.databases.cache.models import SessionQAEntry
+from cognee.infrastructure.databases.cache.models import SessionAgentTraceEntry, SessionQAEntry
 from cognee.infrastructure.databases.exceptions.exceptions import (
     CacheConnectionError,
     SessionQAEntryValidationError,
@@ -38,6 +38,10 @@ class FSCacheAdapter(CacheDBInterface):
     def _session_key(user_id: str, session_id: str) -> str:
         return f"agent_sessions:{user_id}:{session_id}"
 
+    @staticmethod
+    def _agent_trace_key(user_id: str, session_id: str) -> str:
+        return f"agent_traces:{user_id}:{session_id}"
+
     @staticmethod
     def _build_qa_entry_dump(
         question: str,
@@ -62,6 +66,31 @@ class FSCacheAdapter(CacheDBInterface):
         )
         return entry.model_dump()
 
+    @staticmethod
+    def _build_agent_trace_entry_dump(
+        trace_id: str,
+        origin_function: str,
+        status: str,
+        memory_query: str = "",
+        memory_context: str = "",
+        method_params: dict | None = None,
+        method_return_value=None,
+        error_message: str = "",
+        session_feedback: str = "",
+    ) -> dict:
+        entry = SessionAgentTraceEntry(
+            trace_id=trace_id,
+            origin_function=origin_function,
+            status=status,
+            memory_query=memory_query,
+            memory_context=memory_context,
+            method_params=method_params or {},
+            method_return_value=method_return_value,
+            error_message=error_message,
+            session_feedback=session_feedback,
+        )
+        return entry.model_dump()
+
     def _load_entries(self, session_key: str) -> list:
         # Evict expired keys so stale sessions don't linger on disk
         self.cache.expire()
@@ -171,9 +200,10 @@ class FSCacheAdapter(CacheDBInterface):
                 used_graph_element_ids=used_graph_element_ids,
                 memify_metadata=memify_metadata,
             )
-            entries = self._load_entries(session_key)
-            entries.append(qa_entry)
-            self._save_entries(session_key, entries)
+            with self.cache.transact():
+                entries = self._load_entries(session_key)
+                entries.append(qa_entry)
+                self._save_entries(session_key, entries)
         except Exception as e:
             error_msg = f"Unexpected error while adding Q&A to diskcache: {str(e)}"
             logger.error(error_msg)
@@ -210,23 +240,24 @@ class FSCacheAdapter(CacheDBInterface):
         """
         try:
             session_key = self._session_key(user_id, session_id)
-            entries = self._load_entries(session_key)
-            idx = self._find_index_by_qa_id(entries, qa_id)
-            if idx is None:
-                return False
-            merged = self._merge_entry_update(
-                entries[idx],
-                question,
-                context,
-                answer,
-                feedback_text,
-                feedback_score,
-                used_graph_element_ids=used_graph_element_ids,
-                memify_metadata=memify_metadata,
-            )
-            entries[idx] = self._validate_entry_dict(merged)
-            self._save_entries(session_key, entries)
-            return True
+            with self.cache.transact():
+                entries = self._load_entries(session_key)
+                idx = self._find_index_by_qa_id(entries, qa_id)
+                if idx is None:
+                    return False
+                merged = self._merge_entry_update(
+                    entries[idx],
+                    question,
+                    context,
+                    answer,
+                    feedback_text,
+                    feedback_score,
+                    used_graph_element_ids=used_graph_element_ids,
+                    memify_metadata=memify_metadata,
+                )
+                entries[idx] = self._validate_entry_dict(merged)
+                self._save_entries(session_key, entries)
+                return True
         except SessionQAEntryValidationError:
             raise
         except Exception as e:
@@ -240,14 +271,15 @@ class FSCacheAdapter(CacheDBInterface):
         """
         try:
             session_key = self._session_key(user_id, session_id)
-            entries = self._load_entries(session_key)
-            idx = self._find_index_by_qa_id(entries, qa_id)
-            if idx is None:
-                return False
-            merged = self._merge_entry_clear_feedback(entries[idx])
-            entries[idx] = self._validate_entry_dict(merged)
-            self._save_entries(session_key, entries)
-            return True
+            with self.cache.transact():
+                entries = self._load_entries(session_key)
+                idx = self._find_index_by_qa_id(entries, qa_id)
+                if idx is None:
+                    return False
+                merged = self._merge_entry_clear_feedback(entries[idx])
+                entries[idx] = self._validate_entry_dict(merged)
+                self._save_entries(session_key, entries)
+                return True
         except SessionQAEntryValidationError:
             raise
         except Exception as e:
@@ -262,13 +294,14 @@ class FSCacheAdapter(CacheDBInterface):
         """
         try:
             session_key = self._session_key(user_id, session_id)
-            entries = self._load_entries(session_key)
-            idx = self._find_index_by_qa_id(entries, qa_id)
-            if idx is None:
-                return False
-            entries.pop(idx)
-            self._save_entries(session_key, entries)
-            return True
+            with self.cache.transact():
+                entries = self._load_entries(session_key)
+                idx = self._find_index_by_qa_id(entries, qa_id)
+                if idx is None:
+                    return False
+                entries.pop(idx)
+                self._save_entries(session_key, entries)
+                return True
         except Exception as e:
             error_msg = f"Unexpected error while deleting Q&A from diskcache: {str(e)}"
             logger.error(error_msg)
@@ -276,21 +309,72 @@ class FSCacheAdapter(CacheDBInterface):
 
     async def delete_session(self, user_id: str, session_id: str) -> bool:
         """
-        Delete the entire session and all its QA entries.
-        Returns True if deleted, False if session did not exist.
+        Delete the entire session and all its session-scoped artifacts.
+        Returns True if any session data existed, False otherwise.
         """
         try:
             session_key = self._session_key(user_id, session_id)
-            existed = self.cache.get(session_key) is not None
-            if existed:
+            trace_key = self._agent_trace_key(user_id, session_id)
+            qa_existed = self.cache.get(session_key) is not None
+            trace_existed = self.cache.get(trace_key) is not None
+            if qa_existed:
                 self.cache.delete(session_key)
-            return existed
+            if trace_existed:
+                self.cache.delete(trace_key)
+            return qa_existed or trace_existed
 
         except Exception as e:
             error_msg = f"Unexpected error while deleting session from diskcache: {str(e)}"
             logger.error(error_msg)
             raise CacheConnectionError(error_msg) from e
 
+    async def append_agent_trace_step(
+        self,
+        user_id: str,
+        session_id: str,
+        trace_id: str,
+        origin_function: str,
+        status: str,
+        memory_query: str = "",
+        memory_context: str = "",
+        method_params: dict | None = None,
+        method_return_value=None,
+        error_message: str = "",
+        session_feedback: str = "",
+    ) -> None:
+        """Append one trace step to the stored trace list for this session."""
+        try:
+            trace_key = self._agent_trace_key(user_id, session_id)
+            trace_entry = self._build_agent_trace_entry_dump(
+                trace_id=trace_id,
+                origin_function=origin_function,
+                status=status,
+                memory_query=memory_query,
+                memory_context=memory_context,
+                method_params=method_params,
+                method_return_value=method_return_value,
+                error_message=error_message,
+                session_feedback=session_feedback,
+            )
+            with self.cache.transact():
+                entries = self._load_entries(trace_key)
+                entries.append(trace_entry)
+                self._save_entries(trace_key, entries)
+        except Exception as e:
+            error_msg = f"Unexpected error while appending agent trace step to diskcache: {str(e)}"
+            logger.error(error_msg)
+            raise CacheConnectionError(error_msg) from e
+
+    async def get_agent_trace_session(self, user_id: str, session_id: str) -> list[dict]:
+        """Retrieve all stored trace steps for the given session."""
+        trace_key = self._agent_trace_key(user_id, session_id)
+        return self._load_entries(trace_key)
+
+    async def get_agent_trace_feedback(self, user_id: str, session_id: str) -> list[str]:
+        """Retrieve ordered per-step feedback for the given trace session."""
+        entries = await self.get_agent_trace_session(user_id, session_id)
+        return [entry.get("session_feedback", "") for entry in entries]
+
     async def prune(self) -> None:
         """
         Remove all items from the cache. In Cognee, prune means emptying the cache.

{cognee-1.0.1.dev2 → cognee-1.0.1.dev3}/cognee/infrastructure/databases/cache/models.py

@@ -1,5 +1,10 @@
-from typing import Dict, List, Optional
-from pydantic import BaseModel, field_validator
+from typing import Any, Dict, List, Optional
+from pydantic import BaseModel, Field, field_validator
+from cognee.modules.agent_memory.sanitization import (
+    MAX_SERIALIZED_VALUE_LENGTH,
+    sanitize_value,
+    truncate_text,
+)
 
 
 def _validate_list_of_str(value: object, key: str) -> List[str]:
@@ -78,3 +83,62 @@ class SessionQAEntry(BaseModel):
                 raise ValueError("memify_metadata may only have string keys and bool values")
             out[key] = val
         return out if out else None
+
+
+class SessionAgentTraceEntry(BaseModel):
+    """
+    Canonical format for one agent trace step stored in session cache.
+
+    Fields:
+        trace_id: Unique identifier for the trace step.
+        origin_function: Agent method/function that produced the step.
+        status: Execution status for the step.
+        memory_query: Optional memory lookup query used by the step.
+        memory_context: Optional memory context returned to the step.
+        method_params: Serialized method parameters for the step.
+        method_return_value: Serialized method return value for the step.
+        error_message: Error details for failed steps.
+        session_feedback: Per-step feedback generated from the step only.
+    """
+
+    trace_id: str
+    origin_function: str
+    status: str
+    memory_query: str = ""
+    memory_context: str = ""
+    method_params: Dict[str, Any] = Field(default_factory=dict)
+    method_return_value: Any = None
+    error_message: str = ""
+    session_feedback: str = ""
+
+    @field_validator("trace_id", "origin_function", "status")
+    @classmethod
+    def required_non_empty_string(cls, v: str) -> str:
+        if not isinstance(v, str):
+            raise ValueError("value must be a string")
+        stripped = v.strip()
+        if not stripped:
+            raise ValueError("value must be a non-empty string")
+        return truncate_text(stripped, MAX_SERIALIZED_VALUE_LENGTH)
+
+    @field_validator("memory_query", "memory_context", "error_message", "session_feedback")
+    @classmethod
+    def normalize_optional_text_fields(cls, v: str) -> str:
+        if not isinstance(v, str):
+            raise ValueError("value must be a string")
+        return truncate_text(v.strip(), MAX_SERIALIZED_VALUE_LENGTH)
+
+    @field_validator("method_params")
+    @classmethod
+    def sanitize_method_params(cls, v: Dict[str, Any]) -> Dict[str, Any]:
+        if not isinstance(v, dict):
+            raise ValueError("method_params must be a dict")
+        sanitized = sanitize_value(v)
+        if not isinstance(sanitized, dict):
+            raise ValueError("method_params must sanitize to a dict")
+        return sanitized
+
+    @field_validator("method_return_value")
+    @classmethod
+    def sanitize_method_return_value(cls, v: Any) -> Any:
+        return sanitize_value(v)