Merge branch 'dev' into feature/cog-1036-improving-visualization-endp…

…oint

README.md

@@ -130,7 +130,7 @@ This script will run the default pipeline:
 ```python
 import cognee
 import asyncio
-from cognee.api.v1.search import SearchType
+from cognee.modules.search.types import SearchType
 
 async def main():
     # Create a clean slate for cognee -- reset data and system state

cognee-mcp/pyproject.toml

@@ -6,8 +6,8 @@ readme = "README.md"
 requires-python = ">=3.10"
 
 dependencies = [
-    "cognee",
-    "mcp==1.2.0",
+    "cognee[codegraph]",
+    "mcp==1.1.3",
 ]
 
 [[project.authors]]

cognee-mcp/src/server.py

@@ -1,4 +1,5 @@
 import os
+
 import cognee
 import logging
 import importlib.util
@@ -8,7 +9,8 @@
 import mcp.types as types
 from mcp.server import Server, NotificationOptions
 from mcp.server.models import InitializationOptions
-from cognee.api.v1.search import SearchType
+from cognee.api.v1.cognify.code_graph_pipeline import run_code_graph_pipeline
+from cognee.modules.search.types import SearchType
 from cognee.shared.data_models import KnowledgeGraph
 
 mcp = Server("cognee")
@@ -41,6 +43,19 @@ async def list_tools() -> list[types.Tool]:
                 "required": ["text"],
             },
         ),
+        types.Tool(
+            name="codify",
+            description="Transforms codebase into knowledge graph",
+            inputSchema={
+                "type": "object",
+                "properties": {
+                    "repo_path": {
+                        "type": "string",
+                    },
+                },
+                "required": ["repo_path"],
+            },
+        ),
         types.Tool(
             name="search",
             description="Searches for information in knowledge graph",
@@ -51,6 +66,10 @@ async def list_tools() -> list[types.Tool]:
                         "type": "string",
                         "description": "The query to search for",
                     },
+                    "search_type": {
+                        "type": "string",
+                        "description": "The type of search to perform (e.g., INSIGHTS, CODE)",
+                    },
                 },
                 "required": ["search_query"],
             },
@@ -72,15 +91,21 @@ async def call_tools(name: str, arguments: dict) -> list[types.TextContent]:
         with open(os.devnull, "w") as fnull:
             with redirect_stdout(fnull), redirect_stderr(fnull):
                 if name == "cognify":
-                    await cognify(
+                    cognify(
                         text=arguments["text"],
                         graph_model_file=arguments.get("graph_model_file", None),
                         graph_model_name=arguments.get("graph_model_name", None),
                     )
 
                     return [types.TextContent(type="text", text="Ingested")]
+                if name == "codify":
+                    await codify(arguments.get("repo_path"))
+
+                    return [types.TextContent(type="text", text="Indexed")]
                 elif name == "search":
-                    search_results = await search(arguments["search_query"])
+                    search_results = await search(
+                        arguments["search_query"], arguments["search_type"]
+                    )
 
                     return [types.TextContent(type="text", text=search_results)]
                 elif name == "prune":
@@ -102,21 +127,28 @@ async def cognify(text: str, graph_model_file: str = None, graph_model_name: str
     await cognee.add(text)
 
     try:
-        await cognee.cognify(graph_model=graph_model)
+        asyncio.create_task(cognee.cognify(graph_model=graph_model))
     except Exception as e:
         raise ValueError(f"Failed to cognify: {str(e)}")
 
 
-async def search(search_query: str) -> str:
+async def codify(repo_path: str):
+    async for result in run_code_graph_pipeline(repo_path, False):
+        print(result)
+
+
+async def search(search_query: str, search_type: str) -> str:
     """Search the knowledge graph"""
-    search_results = await cognee.search(SearchType.INSIGHTS, query_text=search_query)
+    search_results = await cognee.search(
+        query_type=SearchType[search_type.upper()], query_text=search_query
+    )
 
     results = retrieved_edges_to_string(search_results)
 
     return results
 
 
-async def prune() -> str:
+async def prune():
     """Reset the knowledge graph"""
     await cognee.prune.prune_data()
     await cognee.prune.prune_system(metadata=True)

cognee/api/v1/cognify/cognify_v2.py

@@ -165,7 +165,7 @@ async def get_default_tasks(
                 task_config={"batch_size": 10},
             ),
             Task(add_data_points, task_config={"batch_size": 10}),
-            Task(store_descriptive_metrics),
+            Task(store_descriptive_metrics, include_optional=True),
         ]
     except Exception as error:
         send_telemetry("cognee.cognify DEFAULT TASKS CREATION ERRORED", user.id)

cognee/api/v1/cognify/routers/get_code_pipeline_router.py

@@ -1,20 +1,20 @@
 from fastapi import APIRouter
-from pydantic import BaseModel
+from cognee.api.DTO import InDTO
 from cognee.api.v1.cognify.code_graph_pipeline import run_code_graph_pipeline
 from cognee.modules.retrieval.description_to_codepart_search import (
     code_description_to_code_part_search,
 )
 from fastapi.responses import JSONResponse
 
 
-class CodePipelineIndexPayloadDTO(BaseModel):
+class CodePipelineIndexPayloadDTO(InDTO):
     repo_path: str
     include_docs: bool = False
 
 
-class CodePipelineRetrievePayloadDTO(BaseModel):
+class CodePipelineRetrievePayloadDTO(InDTO):
     query: str
-    fullInput: str
+    full_input: str
 
 
 def get_code_pipeline_router() -> APIRouter:
@@ -34,9 +34,9 @@ async def code_pipeline_retrieve(payload: CodePipelineRetrievePayloadDTO):
         """This endpoint is responsible for retrieving the context."""
         try:
             query = (
-                payload.fullInput.replace("cognee ", "")
-                if payload.fullInput.startswith("cognee ")
-                else payload.fullInput
+                payload.full_input.replace("cognee ", "")
+                if payload.full_input.startswith("cognee ")
+                else payload.full_input
             )
 
             retrieved_codeparts, __ = await code_description_to_code_part_search(
@@ -45,8 +45,8 @@ async def code_pipeline_retrieve(payload: CodePipelineRetrievePayloadDTO):
 
             return [
                 {
-                    "name": codepart.attributes["id"],
-                    "description": codepart.attributes["id"],
+                    "name": codepart.attributes["file_path"],
+                    "description": codepart.attributes["file_path"],
                     "content": codepart.attributes["source_code"],
                 }
                 for codepart in retrieved_codeparts

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

@@ -2,7 +2,7 @@
 from datetime import datetime
 from fastapi import Depends, APIRouter
 from fastapi.responses import JSONResponse
-from cognee.api.v1.search import SearchType
+from cognee.modules.search.types import SearchType
 from cognee.api.DTO import InDTO, OutDTO
 from cognee.modules.users.models import User
 from cognee.modules.search.operations import get_history