cognee 0.2.4__py3-none-any.whl → 0.3.0.dev0__py3-none-any.whl

cognee/api/v1/notebooks/routers/get_notebooks_router.py

@@ -0,0 +1,96 @@
+from uuid import UUID
+from fastapi.encoders import jsonable_encoder
+from fastapi.responses import JSONResponse
+from pydantic import Field
+from typing import List, Optional
+from fastapi import APIRouter, Depends
+
+from cognee.api.DTO import InDTO
+from cognee.infrastructure.databases.relational import get_async_session
+from cognee.infrastructure.utils.run_async import run_async
+from cognee.modules.notebooks.models import Notebook, NotebookCell
+from cognee.modules.notebooks.operations import run_in_local_sandbox
+from cognee.modules.users.models import User
+from cognee.modules.users.methods import get_authenticated_user
+from cognee.modules.notebooks.methods import (
+    create_notebook,
+    delete_notebook,
+    get_notebook,
+    get_notebooks,
+    update_notebook,
+)
+
+
+class NotebookData(InDTO):
+    name: Optional[str] = Field(...)
+    cells: Optional[List[NotebookCell]] = Field(default=[])
+
+
+def get_notebooks_router():
+    router = APIRouter()
+
+    @router.get("")
+    async def get_notebooks_endpoint(user: User = Depends(get_authenticated_user)):
+        return await get_notebooks(user.id)
+
+    @router.post("")
+    async def create_notebook_endpoint(
+        notebook_data: NotebookData, user: User = Depends(get_authenticated_user)
+    ):
+        return await create_notebook(
+            user.id, notebook_data.name, notebook_data.cells, deletable=True
+        )
+
+    @router.put("/{notebook_id}")
+    async def update_notebook_endpoint(
+        notebook_id: UUID, notebook_data: NotebookData, user: User = Depends(get_authenticated_user)
+    ):
+        async with get_async_session(auto_commit=True) as session:
+            notebook: Notebook = await get_notebook(notebook_id, user.id, session)
+
+            if notebook is None:
+                return JSONResponse(status_code=404, content={"error": "Notebook not found"})
+
+            if notebook_data.name and notebook_data.name != notebook.name:
+                notebook.name = notebook_data.name
+
+            if notebook_data.cells:
+                notebook.cells = notebook_data.cells
+
+            return await update_notebook(notebook, session)
+
+    class RunCodeData(InDTO):
+        content: str = Field(...)
+
+    @router.post("/{notebook_id}/{cell_id}/run")
+    async def run_notebook_cell_endpoint(
+        notebook_id: UUID,
+        cell_id: UUID,
+        run_code: RunCodeData,
+        user: User = Depends(get_authenticated_user),
+    ):
+        async with get_async_session() as session:
+            notebook: Notebook = await get_notebook(notebook_id, user.id, session)
+
+            if notebook is None:
+                return JSONResponse(status_code=404, content={"error": "Notebook not found"})
+
+            result, error = await run_async(run_in_local_sandbox, run_code.content)
+
+            return JSONResponse(
+                status_code=200, content={"result": jsonable_encoder(result), "error": error}
+            )
+
+    @router.delete("/{notebook_id}")
+    async def delete_notebook_endpoint(
+        notebook_id: UUID, user: User = Depends(get_authenticated_user)
+    ):
+        async with get_async_session(auto_commit=True) as session:
+            notebook: Notebook = await get_notebook(notebook_id, user.id, session)
+
+            if notebook is None:
+                return JSONResponse(status_code=404, content={"error": "Notebook not found"})
+
+            return await delete_notebook(notebook, session)
+
+    return router

cognee/api/v1/search/routers/get_search_router.py

@@ -4,6 +4,8 @@ from datetime import datetime
 from pydantic import Field
 from fastapi import Depends, APIRouter
 from fastapi.responses import JSONResponse
+from fastapi.encoders import jsonable_encoder
+
 from cognee.modules.search.types import SearchType
 from cognee.api.DTO import InDTO, OutDTO
 from cognee.modules.users.exceptions.exceptions import PermissionDeniedError
@@ -20,7 +22,13 @@ class SearchPayloadDTO(InDTO):
     datasets: Optional[list[str]] = Field(default=None)
     dataset_ids: Optional[list[UUID]] = Field(default=None, examples=[[]])
     query: str = Field(default="What is in the document?")
+    system_prompt: Optional[str] = Field(
+        default="Answer the question using the provided context. Be as brief as possible."
+    )
+    node_name: Optional[list[str]] = Field(default=None, example=[])
     top_k: Optional[int] = Field(default=10)
+    only_context: bool = Field(default=False)
+    use_combined_context: bool = Field(default=False)
 
 
 def get_search_router() -> APIRouter:
@@ -79,7 +87,10 @@ def get_search_router() -> APIRouter:
         - **datasets** (Optional[List[str]]): List of dataset names to search within
         - **dataset_ids** (Optional[List[UUID]]): List of dataset UUIDs to search within
         - **query** (str): The search query string
+        - **system_prompt** Optional[str]: System prompt to be used for Completion type searches in Cognee
+        - **node_name** Optional[list[str]]: Filter results to specific node_sets defined in the add pipeline (for targeted search).
         - **top_k** (Optional[int]): Maximum number of results to return (default: 10)
+        - **only_context** bool: Set to true to only return context Cognee will be sending to LLM in Completion type searches. This will be returned instead of LLM calls for completion type searches.
 
         ## Response
         Returns a list of search results containing relevant nodes from the graph.
@@ -102,7 +113,11 @@ def get_search_router() -> APIRouter:
                 "datasets": payload.datasets,
                 "dataset_ids": [str(dataset_id) for dataset_id in payload.dataset_ids or []],
                 "query": payload.query,
+                "system_prompt": payload.system_prompt,
+                "node_name": payload.node_name,
                 "top_k": payload.top_k,
+                "only_context": payload.only_context,
+                "use_combined_context": payload.use_combined_context,
             },
         )
 
@@ -115,10 +130,14 @@ def get_search_router() -> APIRouter:
                 user=user,
                 datasets=payload.datasets,
                 dataset_ids=payload.dataset_ids,
+                system_prompt=payload.system_prompt,
+                node_name=payload.node_name,
                 top_k=payload.top_k,
+                only_context=payload.only_context,
+                use_combined_context=payload.use_combined_context,
             )
 
-            return results
+            return jsonable_encoder(results)
         except PermissionDeniedError:
             return []
         except Exception as error:

cognee/api/v1/search/search.py

@@ -1,8 +1,9 @@
 from uuid import UUID
 from typing import Union, Optional, List, Type
 
+from cognee.modules.engine.models.node_set import NodeSet
 from cognee.modules.users.models import User
-from cognee.modules.search.types import SearchType
+from cognee.modules.search.types import SearchResult, SearchType, CombinedSearchResult
 from cognee.modules.users.methods import get_default_user
 from cognee.modules.search.methods import search as search_function
 from cognee.modules.data.methods import get_authorized_existing_datasets
@@ -12,16 +13,19 @@ from cognee.modules.data.exceptions import DatasetNotFoundError
 async def search(
     query_text: str,
     query_type: SearchType = SearchType.GRAPH_COMPLETION,
-    user: User = None,
+    user: Optional[User] = None,
     datasets: Optional[Union[list[str], str]] = None,
     dataset_ids: Optional[Union[list[UUID], UUID]] = None,
     system_prompt_path: str = "answer_simple_question.txt",
+    system_prompt: Optional[str] = None,
     top_k: int = 10,
-    node_type: Optional[Type] = None,
+    node_type: Optional[Type] = NodeSet,
     node_name: Optional[List[str]] = None,
     save_interaction: bool = False,
     last_k: Optional[int] = None,
-) -> list:
+    only_context: bool = False,
+    use_combined_context: bool = False,
+) -> Union[List[SearchResult], CombinedSearchResult]:
     """
     Search and query the knowledge graph for insights, information, and connections.
 
@@ -183,11 +187,14 @@ async def search(
         dataset_ids=dataset_ids if dataset_ids else datasets,
         user=user,
         system_prompt_path=system_prompt_path,
+        system_prompt=system_prompt,
         top_k=top_k,
         node_type=node_type,
         node_name=node_name,
         save_interaction=save_interaction,
         last_k=last_k,
+        only_context=only_context,
+        use_combined_context=use_combined_context,
     )
 
     return filtered_search_results

cognee/api/v1/sync/__init__.py

@@ -0,0 +1,17 @@
+from .sync import (
+    sync,
+    SyncResponse,
+    LocalFileInfo,
+    CheckMissingHashesRequest,
+    CheckHashesDiffResponse,
+    PruneDatasetRequest,
+)
+
+__all__ = [
+    "sync",
+    "SyncResponse",
+    "LocalFileInfo",
+    "CheckMissingHashesRequest",
+    "CheckHashesDiffResponse",
+    "PruneDatasetRequest",
+]

cognee/api/v1/sync/routers/__init__.py

@@ -0,0 +1,3 @@
+from .get_sync_router import get_sync_router
+
+__all__ = ["get_sync_router"]

cognee/api/v1/sync/routers/get_sync_router.py

@@ -0,0 +1,241 @@
+from uuid import UUID
+from typing import Optional, List
+from fastapi import APIRouter, Depends
+from fastapi.responses import JSONResponse
+
+
+from cognee.api.DTO import InDTO
+from cognee.modules.users.models import User
+from cognee.modules.users.methods import get_authenticated_user
+from cognee.modules.users.permissions.methods import get_specific_user_permission_datasets
+from cognee.modules.sync.methods import get_running_sync_operations_for_user, get_sync_operation
+from cognee.shared.utils import send_telemetry
+from cognee.shared.logging_utils import get_logger
+from cognee.api.v1.sync import SyncResponse
+from cognee.context_global_variables import set_database_global_context_variables
+
+logger = get_logger()
+
+
+class SyncRequest(InDTO):
+    """Request model for sync operations."""
+
+    dataset_ids: Optional[List[UUID]] = None
+
+
+def get_sync_router() -> APIRouter:
+    router = APIRouter()
+
+    @router.post("", response_model=dict[str, SyncResponse])
+    async def sync_to_cloud(
+        request: SyncRequest,
+        user: User = Depends(get_authenticated_user),
+    ):
+        """
+        Sync local data to Cognee Cloud.
+
+        This endpoint triggers synchronization of local Cognee data to your cloud instance.
+        It uploads your local datasets, knowledge graphs, and processed data to the cloud
+        for backup, sharing, or cloud-based processing.
+
+        ## Request Body (JSON)
+        ```json
+        {
+            "dataset_ids": ["123e4567-e89b-12d3-a456-426614174000", "456e7890-e12b-34c5-d678-901234567000"]
+        }
+        ```
+
+        ## Response
+        Returns immediate response for the sync operation:
+        - **run_id**: Unique identifier for tracking the background sync operation
+        - **status**: Always "started" (operation runs in background)
+        - **dataset_ids**: List of dataset IDs being synced
+        - **dataset_names**: List of dataset names being synced
+        - **message**: Description of the background operation
+        - **timestamp**: When the sync was initiated
+        - **user_id**: User who initiated the sync
+
+        ## Cloud Sync Features
+        - **Automatic Authentication**: Uses your Cognee Cloud credentials
+        - **Data Compression**: Optimizes transfer size for faster uploads
+        - **Smart Sync**: Automatically handles data updates efficiently
+        - **Progress Tracking**: Monitor sync status with sync_id
+        - **Error Recovery**: Automatic retry for failed transfers
+        - **Data Validation**: Ensures data integrity during transfer
+
+        ## Example Usage
+        ```bash
+        # Sync multiple datasets to cloud by IDs (JSON request)
+        curl -X POST "http://localhost:8000/api/v1/sync" \\
+          -H "Content-Type: application/json" \\
+          -H "Cookie: auth_token=your-token" \\
+          -d '{"dataset_ids": ["123e4567-e89b-12d3-a456-426614174000", "456e7890-e12b-34c5-d678-901234567000"]}'
+        
+        # Sync all user datasets (empty request body or null dataset_ids)
+        curl -X POST "http://localhost:8000/api/v1/sync" \\
+          -H "Content-Type: application/json" \\
+          -H "Cookie: auth_token=your-token" \\
+          -d '{}'
+        ```
+
+        ## Error Codes
+        - **400 Bad Request**: Invalid dataset_ids format
+        - **401 Unauthorized**: Invalid or missing authentication
+        - **403 Forbidden**: User doesn't have permission to access dataset
+        - **404 Not Found**: Dataset not found
+        - **409 Conflict**: Sync operation conflict or cloud service unavailable
+        - **413 Payload Too Large**: Dataset too large for current cloud plan
+        - **429 Too Many Requests**: Rate limit exceeded
+
+        ## Notes  
+        - Sync operations run in the background - you get an immediate response
+        - Use the returned run_id to track progress (status API coming soon)
+        - Large datasets are automatically chunked for efficient transfer
+        - Cloud storage usage counts against your plan limits
+        - The sync will continue even if you close your connection
+        """
+        send_telemetry(
+            "Cloud Sync API Endpoint Invoked",
+            user.id,
+            additional_properties={
+                "endpoint": "POST /v1/sync",
+                "dataset_ids": [str(id) for id in request.dataset_ids]
+                if request.dataset_ids
+                else "*",
+            },
+        )
+
+        from cognee.api.v1.sync import sync as cognee_sync
+
+        try:
+            # Check if user has any running sync operations
+            running_syncs = await get_running_sync_operations_for_user(user.id)
+            if running_syncs:
+                # Return information about the existing sync operation
+                existing_sync = running_syncs[0]  # Get the most recent running sync
+                return JSONResponse(
+                    status_code=409,
+                    content={
+                        "error": "Sync operation already in progress",
+                        "details": {
+                            "run_id": existing_sync.run_id,
+                            "status": "already_running",
+                            "dataset_ids": existing_sync.dataset_ids,
+                            "dataset_names": existing_sync.dataset_names,
+                            "message": f"You have a sync operation already in progress with run_id '{existing_sync.run_id}'. Use the status endpoint to monitor progress, or wait for it to complete before starting a new sync.",
+                            "timestamp": existing_sync.created_at.isoformat(),
+                            "progress_percentage": existing_sync.progress_percentage,
+                        },
+                    },
+                )
+
+            # Retrieve existing dataset and check permissions
+            datasets = await get_specific_user_permission_datasets(
+                user.id, "write", request.dataset_ids if request.dataset_ids else None
+            )
+
+            # Execute new cloud sync operation for all datasets
+            sync_result = await cognee_sync(
+                datasets=datasets,
+                user=user,
+            )
+
+            return sync_result
+
+        except ValueError as e:
+            return JSONResponse(status_code=400, content={"error": str(e)})
+        except PermissionError as e:
+            return JSONResponse(status_code=403, content={"error": str(e)})
+        except ConnectionError as e:
+            return JSONResponse(
+                status_code=409, content={"error": f"Cloud service unavailable: {str(e)}"}
+            )
+        except Exception as e:
+            logger.error(f"Cloud sync operation failed: {str(e)}")
+            return JSONResponse(status_code=409, content={"error": "Cloud sync operation failed."})
+
+    @router.get("/status")
+    async def get_sync_status_overview(
+        user: User = Depends(get_authenticated_user),
+    ):
+        """
+        Check if there are any running sync operations for the current user.
+
+        This endpoint provides a simple check to see if the user has any active sync operations
+        without needing to know specific run IDs.
+
+        ## Response
+        Returns a simple status overview:
+        - **has_running_sync**: Boolean indicating if there are any running syncs
+        - **running_sync_count**: Number of currently running sync operations
+        - **latest_running_sync** (optional): Information about the most recent running sync if any exists
+
+        ## Example Usage
+        ```bash
+        curl -X GET "http://localhost:8000/api/v1/sync/status" \\
+          -H "Cookie: auth_token=your-token"
+        ```
+
+        ## Example Responses
+        
+        **No running syncs:**
+        ```json
+        {
+          "has_running_sync": false,
+          "running_sync_count": 0
+        }
+        ```
+
+        **With running sync:**
+        ```json
+        {
+          "has_running_sync": true,
+          "running_sync_count": 1,
+          "latest_running_sync": {
+            "run_id": "12345678-1234-5678-9012-123456789012",
+            "dataset_name": "My Dataset",
+            "progress_percentage": 45,
+            "created_at": "2025-01-01T00:00:00Z"
+          }
+        }
+        ```
+        """
+        send_telemetry(
+            "Sync Status Overview API Endpoint Invoked",
+            user.id,
+            additional_properties={
+                "endpoint": "GET /v1/sync/status",
+            },
+        )
+
+        try:
+            # Get any running sync operations for the user
+            running_syncs = await get_running_sync_operations_for_user(user.id)
+
+            response = {
+                "has_running_sync": len(running_syncs) > 0,
+                "running_sync_count": len(running_syncs),
+            }
+
+            # If there are running syncs, include info about the latest one
+            if running_syncs:
+                latest_sync = running_syncs[0]  # Already ordered by created_at desc
+                response["latest_running_sync"] = {
+                    "run_id": latest_sync.run_id,
+                    "dataset_ids": latest_sync.dataset_ids,
+                    "dataset_names": latest_sync.dataset_names,
+                    "progress_percentage": latest_sync.progress_percentage,
+                    "created_at": latest_sync.created_at.isoformat()
+                    if latest_sync.created_at
+                    else None,
+                }
+
+            return response
+
+        except Exception as e:
+            logger.error(f"Failed to get sync status overview: {str(e)}")
+            return JSONResponse(
+                status_code=500, content={"error": "Failed to get sync status overview"}
+            )
+
+    return router