gmakstutis/cognee
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

fix imports

Browse source
This commit is contained in:
Daulet Amirkhanov 2025-10-31 12:57:33 +00:00
parent 6569dade33
commit e67b6a8bfa
4 changed files with 14 additions and 12 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

6
cognee-mcp/src/tools/clear_registry.py
View file

@ -1,6 +1,6 @@
from contextlib import redirect_stdout from contextlib import redirect_stdout
import sys import sys
from src.utils.context import cognee_client from src.utils import context
from cognee.shared.logging_utils import get_logger from cognee.shared.logging_utils import get_logger
import mcp.types as types import mcp.types as types
@ -21,8 +21,8 @@ async def clear_registry() -> list:
with redirect_stdout(sys.stderr): with redirect_stdout(sys.stderr):
try: try:
await cognee_client.prune_data() await context.cognee_client.prune_data()
await cognee_client.prune_system(metadata=True) await context.cognee_client.prune_system(metadata=True)
logger.info("MCP server registry cleared") logger.info("MCP server registry cleared")
return [ return [
types.TextContent( types.TextContent(

6
cognee-mcp/src/tools/find_mcp_server.py
View file

@ -4,7 +4,7 @@ import sys
from cognee.shared.logging_utils import get_logger from cognee.shared.logging_utils import get_logger
import mcp.types as types import mcp.types as types
from src.utils.context import cognee_client from src.utils import context
logger = get_logger() logger = get_logger()
@ -51,13 +51,13 @@ async def find_mcp_server(requirements: str, max_results: int = 5) -> list:
logger.info(f"Searching for MCP servers matching: {requirements}") logger.info(f"Searching for MCP servers matching: {requirements}")
# Search using GRAPH_COMPLETION for intelligent matching # Search using GRAPH_COMPLETION for intelligent matching
search_results = await cognee_client.search( search_results = await context.cognee_client.search(
query_text=f"Find MCP servers that can: {requirements}. Include their capabilities, installation instructions, and documentation.", query_text=f"Find MCP servers that can: {requirements}. Include their capabilities, installation instructions, and documentation.",
query_type="GRAPH_COMPLETION", query_type="GRAPH_COMPLETION",
) )
# Format the results # Format the results
if cognee_client.use_api: if context.cognee_client.use_api:
if isinstance(search_results, str): if isinstance(search_results, str):
result_text = search_results result_text = search_results
elif isinstance(search_results, list) and len(search_results) > 0: elif isinstance(search_results, list) and len(search_results) > 0:

6
cognee-mcp/src/tools/list_mcp_servers.py
View file

@ -4,7 +4,7 @@ import sys
from cognee.shared.logging_utils import get_logger from cognee.shared.logging_utils import get_logger
import mcp.types as types import mcp.types as types
from src.utils.context import cognee_client from src.utils import context
logger = get_logger() logger = get_logger()
@ -28,13 +28,13 @@ async def list_mcp_servers() -> list:
logger.info("Listing all MCP servers") logger.info("Listing all MCP servers")
# Search for all MCP servers with connection details # Search for all MCP servers with connection details
search_results = await cognee_client.search( search_results = await context.cognee_client.search(
query_text="List all MCP servers with their names, descriptions, capabilities, connection information (URL, command, args), installation instructions, and documentation links", query_text="List all MCP servers with their names, descriptions, capabilities, connection information (URL, command, args), installation instructions, and documentation links",
query_type="GRAPH_COMPLETION", query_type="GRAPH_COMPLETION",
) )
# Format the results # Format the results
if cognee_client.use_api: if context.cognee_client.use_api:
if isinstance(search_results, str): if isinstance(search_results, str):
result_text = search_results result_text = search_results
elif isinstance(search_results, list) and len(search_results) > 0: elif isinstance(search_results, list) and len(search_results) > 0:

8
cognee-mcp/src/tools/remember_mcp_server.py
View file

@ -4,7 +4,7 @@ import sys
from cognee.shared.logging_utils import get_logger, get_log_file_location from cognee.shared.logging_utils import get_logger, get_log_file_location
import mcp.types as types import mcp.types as types
from src.utils.context import cognee_client from src.utils import context
logger = get_logger() logger = get_logger()
@ -115,10 +115,12 @@ Args: {args or "Not provided"}
try: try:
# Add to knowledge graph with special node set # Add to knowledge graph with special node set
await cognee_client.add(server_content, node_set=["mcp_servers", server_name]) await context.cognee_client.add(
server_content, node_set=["mcp_servers", server_name]
)
# Process into knowledge graph # Process into knowledge graph
await cognee_client.cognify() await context.cognee_client.cognify()
logger.info(f"Successfully stored MCP server: {server_name}") logger.info(f"Successfully stored MCP server: {server_name}")
except Exception as e: except Exception as e: