Merge branch 'dev' into test/metrics_in_adapters

Author: alekszievr

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,3 +1,4 @@
+import json
 import os
 import cognee
 import logging
@@ -8,8 +9,10 @@
 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
+from cognee.modules.storage.utils import JSONEncoder
 
 mcp = Server("cognee")
 
@@ -41,6 +44,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 +67,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 +92,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 +128,30 @@ 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:
-    """Search the knowledge graph"""
-    search_results = await cognee.search(SearchType.INSIGHTS, query_text=search_query)
+async def codify(repo_path: str):
+    async for result in run_code_graph_pipeline(repo_path, False):
+        print(result)
 
-    results = retrieved_edges_to_string(search_results)
 
-    return results
+async def search(search_query: str, search_type: str) -> str:
+    """Search the knowledge graph"""
+    search_results = await cognee.search(
+        query_type=SearchType[search_type.upper()], query_text=search_query
+    )
+
+    if search_type.upper() == "CODE":
+        return json.dumps(search_results, cls=JSONEncoder)
+    else:
+        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-mcp/uv.lock

@@ -188,6 +188,7 @@ def start_api_server(host: str = "0.0.0.0", port: int = 8000):
     except Exception as e:
         logger.exception(f"Failed to start server: {e}")
         # Here you could add any cleanup code or error recovery code.
+        raise e
 
 
 if __name__ == "__main__":

cognee/api/client.py

@@ -16,9 +16,22 @@ async def add(
     dataset_name: str = "main_dataset",
     user: User = None,
 ):
+    # Create tables for databases
     await create_relational_db_and_tables()
     await create_pgvector_db_and_tables()
 
+    # Initialize first_run attribute if it doesn't exist
+    if not hasattr(add, "first_run"):
+        add.first_run = True
+
+    if add.first_run:
+        from cognee.infrastructure.llm.utils import test_llm_connection, test_embedding_connection
+
+        # Test LLM and Embedding configuration once before running Cognee
+        await test_llm_connection()
+        await test_embedding_connection()
+        add.first_run = False  # Update flag after first run
+
     if user is None:
         user = await get_default_user()
 

cognee/api/v1/add/add_v2.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/cognify/routers/get_code_pipeline_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