cognee 0.2.3.dev1__py3-none-any.whl → 0.3.0.dev0__py3-none-any.whl

cognee/api/v1/datasets/routers/get_datasets_router.py

@@ -5,6 +5,7 @@ from typing import List, Optional
 from typing_extensions import Annotated
 from fastapi import status
 from fastapi import APIRouter
+from fastapi.encoders import jsonable_encoder
 from fastapi import HTTPException, Query, Depends
 from fastapi.responses import JSONResponse, FileResponse
 
@@ -47,6 +48,7 @@ class DataDTO(OutDTO):
     extension: str
     mime_type: str
     raw_data_location: str
+    dataset_id: UUID
 
 
 class GraphNodeDTO(OutDTO):
@@ -114,7 +116,8 @@ def get_datasets_router() -> APIRouter:
 
     @router.post("", response_model=DatasetDTO)
     async def create_new_dataset(
-        dataset_data: DatasetCreationPayload, user: User = Depends(get_authenticated_user)
+        dataset_data: DatasetCreationPayload,
+        user: User = Depends(get_authenticated_user),
     ):
         """
         Create a new dataset or return existing dataset with the same name.
@@ -284,7 +287,7 @@ def get_datasets_router() -> APIRouter:
         - **500 Internal Server Error**: Error retrieving graph data
         """
 
-        graph_data = await get_formatted_graph_data(dataset_id, user.id)
+        graph_data = await get_formatted_graph_data(dataset_id, user)
 
         return graph_data
 
@@ -327,7 +330,7 @@ def get_datasets_router() -> APIRouter:
             },
         )
 
-        from cognee.modules.data.methods import get_dataset_data, get_dataset
+        from cognee.modules.data.methods import get_dataset_data
 
         # Verify user has permission to read dataset
         dataset = await get_authorized_existing_datasets([dataset_id], "read", user)
@@ -338,12 +341,20 @@ def get_datasets_router() -> APIRouter:
                 content=ErrorResponseDTO(f"Dataset ({str(dataset_id)}) not found."),
             )
 
-        dataset_data = await get_dataset_data(dataset_id=dataset[0].id)
+        dataset_id = dataset[0].id
+
+        dataset_data = await get_dataset_data(dataset_id=dataset_id)
 
         if dataset_data is None:
             return []
 
-        return dataset_data
+        return [
+            dict(
+                **jsonable_encoder(data),
+                dataset_id=dataset_id,
+            )
+            for data in dataset_data
+        ]
 
     @router.get("/status", response_model=dict[str, PipelineRunStatus])
     async def get_dataset_status(

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

@@ -0,0 +1 @@
+from .get_memify_router import get_memify_router

cognee/api/v1/memify/routers/get_memify_router.py

@@ -0,0 +1,100 @@
+from uuid import UUID
+
+from fastapi import APIRouter
+from fastapi.responses import JSONResponse
+from fastapi import Depends
+from pydantic import Field
+from typing import List, Optional, Union, Literal
+
+from cognee.api.DTO import InDTO
+from cognee.modules.users.models import User
+from cognee.modules.users.methods import get_authenticated_user
+from cognee.shared.utils import send_telemetry
+from cognee.modules.pipelines.models import PipelineRunErrored
+from cognee.shared.logging_utils import get_logger
+
+logger = get_logger()
+
+
+class MemifyPayloadDTO(InDTO):
+    extraction_tasks: Optional[List[str]] = Field(
+        default=None,
+        examples=[[]],
+    )
+    enrichment_tasks: Optional[List[str]] = Field(default=None, examples=[[]])
+    data: Optional[str] = Field(default="")
+    dataset_name: Optional[str] = Field(default=None)
+    # Note: Literal is needed for Swagger use
+    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)
+
+
+def get_memify_router() -> APIRouter:
+    router = APIRouter()
+
+    @router.post("", response_model=dict)
+    async def memify(payload: MemifyPayloadDTO, user: User = Depends(get_authenticated_user)):
+        """
+        Enrichment pipeline in Cognee, can work with already built graphs. If no data is provided existing knowledge graph will be used as data,
+        custom data can also be provided instead which can be processed with provided extraction and enrichment tasks.
+
+        Provided tasks and data will be arranged to run the Cognee pipeline and execute graph enrichment/creation.
+
+        ## Request Parameters
+        - **extractionTasks** Optional[List[str]]: List of available Cognee Tasks to execute for graph/data extraction.
+        - **enrichmentTasks** Optional[List[str]]: List of available Cognee Tasks to handle enrichment of provided graph/data from extraction tasks.
+        - **data** Optional[List[str]]: The data to ingest. Can be any text data when custom extraction and enrichment tasks are used.
+              Data provided here will be forwarded to the first extraction task in the pipeline as input.
+              If no data is provided the whole graph (or subgraph if node_name/node_type is specified) will be forwarded
+        - **dataset_name** (Optional[str]): Name of the datasets to memify
+        - **dataset_id** (Optional[UUID]): List of UUIDs of an already existing dataset
+        - **node_name** (Optional[List[str]]):  Filter graph to specific named entities (for targeted search). Used when no data is provided.
+        - **run_in_background** (Optional[bool]): Whether to execute processing asynchronously. Defaults to False (blocking).
+
+        Either datasetName or datasetId must be provided.
+
+        ## Response
+        Returns information about the add operation containing:
+        - Status of the operation
+        - Details about the processed data
+        - Any relevant metadata from the ingestion process
+
+        ## Error Codes
+        - **400 Bad Request**: Neither datasetId nor datasetName provided
+        - **409 Conflict**: Error during memify operation
+        - **403 Forbidden**: User doesn't have permission to use dataset
+
+        ## Notes
+        - To memify datasets not owned by the user, use dataset_id (when ENABLE_BACKEND_ACCESS_CONTROL is set to True)
+        - datasetId value can only be the UUID of an already existing dataset
+        """
+
+        send_telemetry(
+            "Memify API Endpoint Invoked",
+            user.id,
+            additional_properties={"endpoint": "POST /v1/memify"},
+        )
+
+        if not payload.dataset_id and not payload.dataset_name:
+            raise ValueError("Either datasetId or datasetName must be provided.")
+
+        try:
+            from cognee.modules.memify import memify as cognee_memify
+
+            memify_run = await cognee_memify(
+                extraction_tasks=payload.extraction_tasks,
+                enrichment_tasks=payload.enrichment_tasks,
+                data=payload.data,
+                dataset=payload.dataset_id if payload.dataset_id else payload.dataset_name,
+                node_name=payload.node_name,
+                user=user,
+            )
+
+            if isinstance(memify_run, PipelineRunErrored):
+                return JSONResponse(status_code=420, content=memify_run)
+            return memify_run
+        except Exception as error:
+            return JSONResponse(status_code=409, content={"error": str(error)})
+
+    return router

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

@@ -0,0 +1 @@
+from .get_notebooks_router import get_notebooks_router

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/responses/default_tools.py

@@ -49,6 +49,10 @@ DEFAULT_TOOLS = [
                     "type": "string",
                     "description": "Path to a custom ontology file",
                 },
+                "custom_prompt": {
+                    "type": "string",
+                    "description": "Custom prompt for entity extraction and graph generation. If provided, this prompt will be used instead of the default prompts.",
+                },
             },
             "required": ["text"],
         },

cognee/api/v1/responses/dispatch_function.py

@@ -88,11 +88,16 @@ async def handle_cognify(arguments: Dict[str, Any], user) -> str:
     """Handle cognify function call"""
     text = arguments.get("text")
     ontology_file_path = arguments.get("ontology_file_path")
+    custom_prompt = arguments.get("custom_prompt")
 
     if text:
         await add(data=text, user=user)
 
-    await cognify(user=user, ontology_file_path=ontology_file_path if ontology_file_path else None)
+    await cognify(
+        user=user,
+        ontology_file_path=ontology_file_path if ontology_file_path else None,
+        custom_prompt=custom_prompt,
+    )
 
     return (
         "Text successfully converted into knowledge graph."

cognee/api/v1/responses/models.py

@@ -70,7 +70,7 @@ class ResponseRequest(InDTO):
     tool_choice: Optional[Union[str, Dict[str, Any]]] = "auto"
     user: Optional[str] = None
     temperature: Optional[float] = 1.0
-    max_tokens: Optional[int] = None
+    max_completion_tokens: Optional[int] = None
 
 
 class ToolCallOutput(BaseModel):

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,14 +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,
-) -> list:
+    save_interaction: bool = False,
+    last_k: Optional[int] = None,
+    only_context: bool = False,
+    use_combined_context: bool = False,
+) -> Union[List[SearchResult], CombinedSearchResult]:
     """
     Search and query the knowledge graph for insights, information, and connections.
 
@@ -107,6 +113,8 @@ async def search(
 
         node_name: Filter results to specific named entities (for targeted search).
 
+        save_interaction: Save interaction (query, context, answer connected to triplet endpoints) results into the graph or not
+
     Returns:
         list: Search results in format determined by query_type:
 
@@ -179,9 +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"]