docs: Add comprehensive docstrings and fix default top_k consistency

Address PR feedback from CodeRabbit AI:
- Add detailed docstring for search_task internal function
- Document top_k parameter in main search function docstring
- Fix default top_k inconsistency (was 10 in client, now 5 everywhere)
- Clarify performance implications of different top_k values

Changes:
- server.py: Add top_k parameter documentation and search_task docstring
- cognee_client.py: Change default top_k from 10 to 5 for consistency

This ensures consistent behavior across the MCP call chain and
provides clear guidance for users on choosing appropriate top_k values.
This commit is contained in:
AnveshJarabani 2026-01-03 01:33:13 -06:00
parent 7ee36f883b
commit 6a5ba70ced
No known key found for this signature in database
GPG key ID: 931E4CB5B2BF7B1D
2 changed files with 28 additions and 2 deletions

View file

@ -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.

View file

@ -389,6 +389,13 @@ async def search(search_query: str, search_type: str, top_k: int = 5) -> 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
@ -426,7 +433,26 @@ async def search(search_query: str, search_type: str, top_k: int = 5) -> list:
"""
async def search_task(search_query: str, search_type: str, top_k: int) -> str:
"""Search the knowledge graph"""
"""
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):