cognee 1.0.1.dev2__tar.gz → 1.0.1.dev4__tar.gz

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: cognee
-Version: 1.0.1.dev2
+Version: 1.0.1.dev4
 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.dev4}/cognee/__init__.py

@@ -46,6 +46,7 @@ from cognee.run_migrations import run_startup_migrations
 # V2 memory-oriented API
 # ---------------------------------------------------------------------------
 from .api.v1 import remember, RememberResult, recall, improve, forget, serve, disconnect, visualize
+from .memory import MemoryEntry, QAEntry, TraceEntry, FeedbackEntry
 
 # Tracing / Observability
 from cognee.modules.observability.trace_context import (

{cognee-1.0.1.dev2 → cognee-1.0.1.dev4}/cognee/api/client.py

@@ -47,6 +47,7 @@ from cognee.api.v1.users.routers import (
 )
 from cognee.api.v1.api_keys.routers import get_api_key_management_router
 from cognee.api.v1.activity.routers import get_activity_router
+from cognee.api.v1.sessions import get_sessions_router
 from cognee.modules.users.methods.get_authenticated_user import REQUIRE_AUTHENTICATION
 
 # Ensure application logging is configured for container stdout/stderr
@@ -292,6 +293,13 @@ app.include_router(
     tags=["activity"],
 )
 
+# Sessions lifecycle + dashboard aggregates
+app.include_router(
+    get_sessions_router(),
+    prefix="/api/v1/sessions",
+    tags=["sessions"],
+)
+
 app.include_router(get_remember_router(), prefix="/api/v1/remember", tags=["remember"])
 app.include_router(get_recall_router(), prefix="/api/v1/recall", tags=["recall"])
 app.include_router(get_improve_router(), prefix="/api/v1/improve", tags=["improve"])

{cognee-1.0.1.dev2 → cognee-1.0.1.dev4}/cognee/api/v1/activity/routers/get_activity_router.py

@@ -9,6 +9,7 @@ from uuid import UUID
 from fastapi import APIRouter, Query, Depends
 from cognee.modules.users.models import User
 from cognee.modules.users.methods.get_authenticated_user import get_authenticated_user
+from cognee.modules.users.exceptions import PermissionDeniedError
 from cognee.modules.users.permissions.methods.get_specific_user_permission_datasets import (
     get_specific_user_permission_datasets,
 )
@@ -28,10 +29,20 @@ def get_activity_router() -> APIRouter:
         from cognee.modules.users.models import User
         from sqlalchemy import select, outerjoin
 
-        dataset_ids = await get_specific_user_permission_datasets(
-            user.id, "read", [dataset_id] if dataset_id else None
-        )
-        dataset_ids = [ds.id for ds in dataset_ids]
+        try:
+            permitted_datasets = await get_specific_user_permission_datasets(
+                user.id, "read", [dataset_id] if dataset_id else None
+            )
+        except PermissionDeniedError:
+            # For list requests, treat "no accessible datasets" as an empty activity feed.
+            # Keep explicit dataset requests strict by propagating the original 403.
+            if dataset_id is None:
+                return []
+            raise
+
+        dataset_ids = [ds.id for ds in permitted_datasets]
+        if not dataset_ids:
+            return []
 
         db_engine = get_relational_engine()
         async with db_engine.get_async_session() as session:

{cognee-1.0.1.dev2 → cognee-1.0.1.dev4}/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.dev4}/cognee/api/v1/improve/improve.py

@@ -119,16 +119,57 @@ async def improve(
 
         feedback_alpha = kwargs.pop("feedback_alpha", 0.1)
 
+        # Mutex: single-session improves serialize on the session's
+        # lock so auto-improve + idle-watcher + SessionEnd don't
+        # duplicate work. Multi-session improves skip the lock — the
+        # pattern is rare and locking N sessions at once is messy.
+        acquired_lock_for: Optional[str] = None
+        if session_ids and len(session_ids) == 1:
+            from cognee.infrastructure.locks import (
+                release_improve_lock,
+                try_acquire_improve_lock,
+            )
+
+            sole_session = session_ids[0]
+            if not await try_acquire_improve_lock(sole_session):
+                logger.info(
+                    "improve: session '%s' already being improved, skipping",
+                    sole_session,
+                )
+                return {}
+            acquired_lock_for = sole_session
+        else:
+            release_improve_lock = None  # type: ignore[assignment]
+
         # Stage 1 & 2: bridge sessions into the permanent graph
         if session_ids:
-            await _bridge_sessions(
-                dataset=dataset,
-                session_ids=session_ids,
-                user=user,
-                feedback_alpha=feedback_alpha,
-                run_in_background=run_in_background,
-            )
-            stages_run.extend(["feedback_weights", "persist_sessions"])
+            try:
+                await _bridge_sessions(
+                    dataset=dataset,
+                    session_ids=session_ids,
+                    user=user,
+                    feedback_alpha=feedback_alpha,
+                    run_in_background=run_in_background,
+                )
+                stages_run.extend(["feedback_weights", "persist_sessions"])
+
+                # Stage 2b: persist agent trace steps (tool calls with
+                # per-step feedback) into the graph. Without this, the
+                # plugin's trace activity never reaches permanent
+                # memory — only QA entries do.
+                await _persist_session_traces(
+                    dataset=dataset,
+                    session_ids=session_ids,
+                    user=user,
+                    run_in_background=run_in_background,
+                )
+                stages_run.append("persist_trace_steps")
+            except Exception:
+                if acquired_lock_for:
+                    from cognee.infrastructure.locks import release_improve_lock
+
+                    await release_improve_lock(acquired_lock_for)
+                raise
 
         # Stage 3: default enrichment (triplet embeddings)
         from cognee.modules.memify import memify
@@ -158,6 +199,12 @@ async def improve(
             stages_run.append("sync_graph_to_sessions")
 
         span.set_attribute(COGNEE_IMPROVE_STAGES, ",".join(stages_run))
+
+        if acquired_lock_for:
+            from cognee.infrastructure.locks import release_improve_lock
+
+            await release_improve_lock(acquired_lock_for)
+
         return result
 
 
@@ -225,6 +272,47 @@ async def _bridge_sessions(
         logger.warning("improve: session persistence failed (non-fatal): %s", e)
 
 
+async def _persist_session_traces(
+    dataset: Union[str, UUID],
+    session_ids: List[str],
+    user,
+    run_in_background: bool,
+):
+    """Cognify per-step agent trace feedbacks into the knowledge graph.
+
+    Without this step, the Claude Code plugin's tool-call activity
+    (the bulk of session data — hundreds of Bash/Edit/Read/Write trace
+    steps per session) never makes it into permanent memory. Only QA
+    entries do via ``persist_sessions_in_knowledge_graph_pipeline``.
+
+    Runs the dedicated ``persist_agent_trace_feedbacks_in_knowledge_graph_pipeline``
+    that extracts per-step ``session_feedback`` from the cache and
+    cognifies it into the ``agent_trace_feedbacks`` node-set.
+    """
+    dataset_name = await _resolve_dataset_name(dataset, user)
+
+    try:
+        from cognee.memify_pipelines.persist_agent_trace_feedbacks_in_knowledge_graph import (
+            persist_agent_trace_feedbacks_in_knowledge_graph_pipeline,
+        )
+
+        await persist_agent_trace_feedbacks_in_knowledge_graph_pipeline(
+            user=user,
+            session_ids=session_ids,
+            dataset=dataset_name,
+            node_set_name="agent_trace_feedbacks",
+            raw_trace_content=False,
+            last_n_steps=None,  # persist all stored steps on demand
+            run_in_background=run_in_background,
+        )
+        logger.info(
+            "improve: agent trace steps persisted from %d session(s)",
+            len(session_ids),
+        )
+    except Exception as e:
+        logger.warning("improve: trace persistence failed (non-fatal): %s", e)
+
+
 async def _sync_graph_to_sessions(
     dataset: Union[str, UUID],
     session_ids: List[str],

{cognee-1.0.1.dev2 → cognee-1.0.1.dev4}/cognee/api/v1/improve/routers/get_improve_router.py

@@ -26,6 +26,11 @@ class ImprovePayloadDTO(InDTO):
     dataset_id: Union[UUID, Literal[""], None] = Field(default=None, examples=[""])
     node_name: Optional[List[str]] = Field(default=None, examples=[[]])
     run_in_background: Optional[bool] = Field(default=False)
+    # Session IDs to bridge into the permanent graph. When set, improve
+    # runs the full session pipeline (feedback weights + QA persist +
+    # trace-step persist + graph→session sync) in addition to the
+    # default memify enrichment.
+    session_ids: Optional[List[str]] = Field(default=None, examples=[[]])
 
 
 def get_improve_router() -> APIRouter:
@@ -79,6 +84,7 @@ def get_improve_router() -> APIRouter:
                 data=payload.data,
                 dataset=payload.dataset_id if payload.dataset_id else payload.dataset_name,
                 node_name=payload.node_name,
+                session_ids=payload.session_ids,
                 user=user,
                 run_in_background=payload.run_in_background,
             )

{cognee-1.0.1.dev2 → cognee-1.0.1.dev4}/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