diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index 6ca815825..87e3dc91c 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -97,6 +97,21 @@ git checkout -b feature/your-feature-name python cognee/cognee/tests/test_library.py ``` +### Running Simple Example + +Change .env.example into .env and provide your OPENAI_API_KEY as LLM_API_KEY + +Make sure to run ```shell uv sync ``` in the root cloned folder or set up a virtual environment to run cognee + +```shell +python cognee/cognee/examples/python/simple_example.py +``` +or + +```shell +uv run python cognee/cognee/examples/python/simple_example.py +``` + ## 4. 📤 Submitting Changes 1. Install ruff on your system diff --git a/cognee-mcp/src/cognee_client.py b/cognee-mcp/src/cognee_client.py index a2fd3345f..9d98cb0b5 100644 --- a/cognee-mcp/src/cognee_client.py +++ b/cognee-mcp/src/cognee_client.py @@ -151,7 +151,7 @@ class CogneeClient: query_type: str, datasets: Optional[List[str]] = None, system_prompt: Optional[str] = None, - top_k: int = 10, + top_k: int = 5, ) -> Any: """ Search the knowledge graph. @@ -192,7 +192,9 @@ class CogneeClient: with redirect_stdout(sys.stderr): results = await self.cognee.search( - query_type=SearchType[query_type.upper()], query_text=query_text + query_type=SearchType[query_type.upper()], + query_text=query_text, + top_k=top_k ) return results diff --git a/cognee-mcp/src/server.py b/cognee-mcp/src/server.py index 01dee6479..f67b62648 100755 --- a/cognee-mcp/src/server.py +++ b/cognee-mcp/src/server.py @@ -316,7 +316,7 @@ async def save_interaction(data: str) -> list: @mcp.tool() -async def search(search_query: str, search_type: str) -> list: +async def search(search_query: str, search_type: str, top_k: int = 5) -> list: """ Search and query the knowledge graph for insights, information, and connections. @@ -389,6 +389,13 @@ async def search(search_query: str, search_type: str) -> list: The search_type is case-insensitive and will be converted to uppercase. + top_k : int, optional + Maximum number of results to return (default: 5). + Controls the amount of context retrieved from the knowledge graph. + - Lower values (3-5): Faster, more focused results + - Higher values (10-20): More comprehensive, but slower and more context-heavy + Helps manage response size and context window usage in MCP clients. + Returns ------- list @@ -425,13 +432,32 @@ async def search(search_query: str, search_type: str) -> list: """ - async def search_task(search_query: str, search_type: str) -> str: - """Search the knowledge graph""" + async def search_task(search_query: str, search_type: str, top_k: int) -> str: + """ + Internal task to execute knowledge graph search with result formatting. + + Handles the actual search execution and formats results appropriately + for MCP clients based on the search type and execution mode (API vs direct). + + Parameters + ---------- + search_query : str + The search query in natural language + search_type : str + Type of search to perform (GRAPH_COMPLETION, CHUNKS, etc.) + top_k : int + Maximum number of results to return + + Returns + ------- + str + Formatted search results as a string, with format depending on search_type + """ # NOTE: MCP uses stdout to communicate, we must redirect all output # going to stdout ( like the print function ) to stderr. with redirect_stdout(sys.stderr): search_results = await cognee_client.search( - query_text=search_query, query_type=search_type + query_text=search_query, query_type=search_type, top_k=top_k ) # Handle different result formats based on API vs direct mode @@ -465,7 +491,7 @@ async def search(search_query: str, search_type: str) -> list: else: return str(search_results) - search_results = await search_task(search_query, search_type) + search_results = await search_task(search_query, search_type, top_k) return [types.TextContent(type="text", text=search_results)]