diff --git a/mcp_server/README.md b/mcp_server/README.md index 9f432aea73..d613fbca7b 100644 --- a/mcp_server/README.md +++ b/mcp_server/README.md @@ -251,6 +251,25 @@ The `config.yaml` file supports environment variable expansion using `${VAR_NAME You can set these variables in a `.env` file in the project directory. +### Search tuning (env vars, prefix `SEARCH__`) + +These knobs tune the `search_memory_facts` and `search_nodes` tools. They are +nested config, so the env prefix is `SEARCH__` (double underscore). + +| Var | Default | Purpose | +|---|---|---| +| `SEARCH__RERANKER` | `rrf` | `rrf` (reciprocal-rank fusion; relevance-ordered and diverse — the reliable default), `mmr` (diversity-aware but miscalibrated for low-magnitude embedding similarities — validate before use), or `cross_encoder`. | +| `SEARCH__MMR_LAMBDA` | `0.5` | MMR relevance/diversity tradeoff; 1.0 = no diversity. Only used when reranker is `mmr`. | +| `SEARCH__RERANKER_MIN_SCORE` | `0.0` | Min score to keep a result. Only meaningful for `cross_encoder`; ignored for `mmr` (which must not filter) and left at 0 for `rrf`. | +| `SEARCH__MAX_FACTS` | `6` | Default result count for `search_memory_facts`. | +| `SEARCH__MAX_NODES` | `6` | Default result count for `search_nodes`. | +| `SEARCH__EXCLUDE_INVALIDATED` | `true` | Exclude superseded/expired facts by default. Per-call override: `include_invalidated=true`. | + +**Cross-encoder caveat:** `SEARCH__RERANKER=cross_encoder` gives a true relevance +floor but requires a cross-encoder client configured on the Graphiti instance and +adds a model call per search (latency + another provider dependency). It is **not** +the default; validate reliability before enabling. + ## Running the Server ### Default Setup (FalkorDB Combined Container) diff --git a/mcp_server/src/config/schema.py b/mcp_server/src/config/schema.py index 6417c73e76..706dc6c355 100644 --- a/mcp_server/src/config/schema.py +++ b/mcp_server/src/config/schema.py @@ -2,7 +2,7 @@ import os from pathlib import Path -from typing import Any +from typing import Any, Literal import yaml from pydantic import BaseModel, Field @@ -246,6 +246,38 @@ class EdgeTypeMapEntry(BaseModel): ) +class SearchTuningConfig(BaseModel): + """Tuning knobs for the MCP search tools (search_nodes / search_memory_facts).""" + + reranker: Literal['rrf', 'mmr', 'cross_encoder'] = Field( + default='rrf', + description="Reranker for hybrid search. 'rrf' is plain reciprocal-rank " + "fusion (the reliable default: relevance-ordered, diverse). 'mmr' penalizes " + 'redundancy but is miscalibrated for low-magnitude embedding similarities ' + '(it can rank outliers over relevant clustered facts); validate before use. ' + "'cross_encoder' scores true relevance but requires a configured " + 'cross-encoder client and adds a model call per search.', + ) + mmr_lambda: float = Field( + default=0.5, + description='MMR relevance/diversity tradeoff (1.0 = pure relevance, no ' + 'diversity; lower = more diversity). Only used when reranker == mmr.', + ) + reranker_min_score: float = Field( + default=0.0, + description='Minimum reranker score to keep a result. Meaningful only for ' + 'cross_encoder (0-1 calibrated); ignored for mmr (which must not filter) ' + 'and left at 0 for rrf.', + ) + max_facts: int = Field(default=6, description='Default max facts for search_memory_facts') + max_nodes: int = Field(default=6, description='Default max nodes for search_nodes') + exclude_invalidated: bool = Field( + default=True, + description='Exclude superseded/expired edges from fact search by default. ' + 'Clients can opt back in per call with include_invalidated=True.', + ) + + class GraphitiAppConfig(BaseModel): """Graphiti-specific configuration.""" @@ -270,6 +302,7 @@ class GraphitiConfig(BaseSettings): embedder: EmbedderConfig = Field(default_factory=EmbedderConfig) database: DatabaseConfig = Field(default_factory=DatabaseConfig) graphiti: GraphitiAppConfig = Field(default_factory=GraphitiAppConfig) + search: SearchTuningConfig = Field(default_factory=SearchTuningConfig) # Additional server options destroy_graph: bool = Field(default=False, description='Clear graph on startup') diff --git a/mcp_server/src/graphiti_mcp_server.py b/mcp_server/src/graphiti_mcp_server.py index efd9536de3..cdfd95ef7d 100644 --- a/mcp_server/src/graphiti_mcp_server.py +++ b/mcp_server/src/graphiti_mcp_server.py @@ -40,6 +40,11 @@ from services.factories import DatabaseDriverFactory, EmbedderFactory, LLMClientFactory from services.queue_service import QueueService from utils.formatting import format_fact_result, to_edge_result, to_node_result +from utils.search_tuning import ( + apply_liveness_filter, + build_edge_search_config, + build_node_search_config, +) from utils.type_config import ( build_edge_type_map, build_edge_types, @@ -478,7 +483,7 @@ async def add_memory( async def search_nodes( query: str, group_ids: str | list[str] | None = None, - max_nodes: int = 10, + max_nodes: int | None = None, entity_types: list[str] | None = None, center_node_uuid: str | None = None, ) -> NodeSearchResponse | ErrorResponse: @@ -488,7 +493,8 @@ async def search_nodes( query: The search query group_ids: Optional group ID, or list of group IDs, to filter results (a single string is accepted and treated as a one-element list) - max_nodes: Maximum number of nodes to return (default: 10) + max_nodes: Maximum number of nodes to return (defaults to the server's + search.max_nodes setting) entity_types: Optional list of entity type names (node labels) to filter by center_node_uuid: Optional UUID of a node to center the search around. Results closer to this node in the graph are ranked higher. @@ -499,6 +505,8 @@ async def search_nodes( return ErrorResponse(error='Graphiti service not initialized') try: + max_nodes = max_nodes if max_nodes is not None else config.search.max_nodes + client = await graphiti_service.get_client() # Accept a scalar group_id or a list; fall back to the default when omitted. @@ -511,22 +519,16 @@ async def search_nodes( else [] ) - # Create search filters - search_filters = SearchFilters( - node_labels=entity_types, - ) + search_filters = SearchFilters(node_labels=entity_types) - # center_node_uuid is only honored by the node_distance reranker, so select - # that recipe when a center node is given (mirroring core's Graphiti.search); - # otherwise use RRF. - from graphiti_core.search.search_config_recipes import ( - NODE_HYBRID_SEARCH_NODE_DISTANCE, - NODE_HYBRID_SEARCH_RRF, + node_config = build_node_search_config( + reranker=config.search.reranker, + mmr_lambda=config.search.mmr_lambda, + limit=max_nodes, + min_score=config.search.reranker_min_score, + center_node_uuid=center_node_uuid, ) - node_config = ( - NODE_HYBRID_SEARCH_NODE_DISTANCE if center_node_uuid else NODE_HYBRID_SEARCH_RRF - ) results = await client.search_( query=query, config=node_config, @@ -535,15 +537,16 @@ async def search_nodes( search_filter=search_filters, ) - # Extract nodes from results nodes = results.nodes[:max_nodes] if results.nodes else [] + scores = results.node_reranker_scores[:max_nodes] if results.node_reranker_scores else [] if not nodes: return NodeSearchResponse(message='No relevant nodes found', nodes=[]) - # Format the results (embeddings stripped by to_node_result) - node_results = [to_node_result(node) for node in nodes] - + node_results = [ + to_node_result(node, score=(scores[i] if i < len(scores) else None)) + for i, node in enumerate(nodes) + ] return NodeSearchResponse(message='Nodes retrieved successfully', nodes=node_results) except Exception as e: error_msg = str(e) @@ -555,13 +558,14 @@ async def search_nodes( async def search_memory_facts( query: str, group_ids: str | list[str] | None = None, - max_facts: int = 10, + max_facts: int | None = None, center_node_uuid: str | None = None, edge_types: list[str] | None = None, valid_at_after: str | None = None, valid_at_before: str | None = None, invalid_at_after: str | None = None, invalid_at_before: str | None = None, + include_invalidated: bool | None = None, ) -> FactSearchResponse | ErrorResponse: """Search the graph memory for relevant facts (entity edges). @@ -569,7 +573,8 @@ async def search_memory_facts( query: The search query group_ids: Optional group ID, or list of group IDs, to filter results (a single string is accepted and treated as a one-element list) - max_facts: Maximum number of facts to return (default: 10) + max_facts: Maximum number of facts to return (defaults to the server's + search.max_facts setting) center_node_uuid: Optional UUID of a node to center the search around edge_types: Optional list of edge (fact) type names to filter by valid_at_after: Optional ISO-8601 lower bound; only facts whose valid_at is at or @@ -577,6 +582,9 @@ async def search_memory_facts( valid_at_before: Optional ISO-8601 upper bound on a fact's valid_at invalid_at_after: Optional ISO-8601 lower bound on a fact's invalid_at invalid_at_before: Optional ISO-8601 upper bound on a fact's invalid_at + include_invalidated: When True, include superseded/expired facts in + results (for historical/temporal reasoning). Defaults to the + server's search.exclude_invalidated setting (excluded by default). """ global graphiti_service @@ -584,7 +592,14 @@ async def search_memory_facts( return ErrorResponse(error='Graphiti service not initialized') try: - # Validate max_facts parameter + # Resolve defaults from config. + max_facts = max_facts if max_facts is not None else config.search.max_facts + include_invalidated = ( + include_invalidated + if include_invalidated is not None + else (not config.search.exclude_invalidated) + ) + if max_facts <= 0: return ErrorResponse(error='max_facts must be a positive integer') @@ -600,6 +615,9 @@ async def search_memory_facts( except ValueError as e: return ErrorResponse(error=f'Invalid date filter: {e}') + if not include_invalidated: + search_filter = apply_liveness_filter(search_filter) + client = await graphiti_service.get_client() # Accept a scalar group_id or a list; fall back to the default when omitted. @@ -612,18 +630,32 @@ async def search_memory_facts( else [] ) - relevant_edges = await client.search( - group_ids=effective_group_ids, + edge_config = build_edge_search_config( + reranker=config.search.reranker, + mmr_lambda=config.search.mmr_lambda, + limit=max_facts, + min_score=config.search.reranker_min_score, + center_node_uuid=center_node_uuid, + ) + + results = await client.search_( query=query, - num_results=max_facts, + config=edge_config, + group_ids=effective_group_ids, center_node_uuid=center_node_uuid, - search_filter=search_filter, + search_filter=search_filter if search_filter is not None else SearchFilters(), ) - if not relevant_edges: + edges = results.edges[:max_facts] if results.edges else [] + scores = results.edge_reranker_scores[:max_facts] if results.edge_reranker_scores else [] + + if not edges: return FactSearchResponse(message='No relevant facts found', facts=[]) - facts = [format_fact_result(edge) for edge in relevant_edges] + facts = [ + format_fact_result(edge, score=(scores[i] if i < len(scores) else None)) + for i, edge in enumerate(edges) + ] return FactSearchResponse(message='Facts retrieved successfully', facts=facts) except Exception as e: error_msg = str(e) diff --git a/mcp_server/src/models/response_types.py b/mcp_server/src/models/response_types.py index 387f9b1061..8408ad4911 100644 --- a/mcp_server/src/models/response_types.py +++ b/mcp_server/src/models/response_types.py @@ -21,6 +21,7 @@ class NodeResult(TypedDict): summary: str | None group_id: str attributes: dict[str, Any] + score: float | None class NodeSearchResponse(TypedDict): @@ -74,6 +75,7 @@ class EdgeResult(TypedDict): created_at: str | None valid_at: str | None invalid_at: str | None + score: float | None class TripletResponse(TypedDict): diff --git a/mcp_server/src/utils/formatting.py b/mcp_server/src/utils/formatting.py index 1b5b7f3b1a..ea169bcab1 100644 --- a/mcp_server/src/utils/formatting.py +++ b/mcp_server/src/utils/formatting.py @@ -8,7 +8,7 @@ from models.response_types import EdgeResult, NodeResult -def to_node_result(node: EntityNode) -> NodeResult: +def to_node_result(node: EntityNode, score: float | None = None) -> NodeResult: """Build a NodeResult TypedDict from an EntityNode, dropping embeddings.""" attrs = node.attributes if node.attributes else {} attrs = {k: v for k, v in attrs.items() if 'embedding' not in k.lower()} @@ -20,10 +20,11 @@ def to_node_result(node: EntityNode) -> NodeResult: summary=node.summary, group_id=node.group_id, attributes=attrs, + score=score, ) -def to_edge_result(edge: EntityEdge) -> EdgeResult: +def to_edge_result(edge: EntityEdge, score: float | None = None) -> EdgeResult: """Build an EdgeResult TypedDict from an EntityEdge.""" return EdgeResult( uuid=edge.uuid, @@ -35,6 +36,7 @@ def to_edge_result(edge: EntityEdge) -> EdgeResult: created_at=edge.created_at.isoformat() if edge.created_at else None, valid_at=edge.valid_at.isoformat() if edge.valid_at else None, invalid_at=edge.invalid_at.isoformat() if edge.invalid_at else None, + score=score, ) @@ -61,13 +63,14 @@ def format_node_result(node: EntityNode) -> dict[str, Any]: return result -def format_fact_result(edge: EntityEdge) -> dict[str, Any]: - """Format an entity edge into a readable result. +def format_fact_result(edge: EntityEdge, score: float | None = None) -> dict[str, Any]: + """Format an entity edge into a readable result, including its rerank score. Since EntityEdge is a Pydantic BaseModel, we can use its built-in serialization capabilities. Args: edge: The EntityEdge to format + score: Optional reranker score for this edge Returns: A dictionary representation of the edge with serialized dates and excluded embeddings @@ -79,4 +82,5 @@ def format_fact_result(edge: EntityEdge) -> dict[str, Any]: }, ) result.get('attributes', {}).pop('fact_embedding', None) + result['score'] = score return result diff --git a/mcp_server/src/utils/search_tuning.py b/mcp_server/src/utils/search_tuning.py new file mode 100644 index 0000000000..e280558964 --- /dev/null +++ b/mcp_server/src/utils/search_tuning.py @@ -0,0 +1,111 @@ +"""Search-config and filter helpers for the MCP search tools. + +Composes graphiti_core primitives without mutating the shared recipe +singletons: every builder returns a deep copy. +""" + +from __future__ import annotations + +from datetime import datetime, timezone + +from graphiti_core.search.search_config import EdgeReranker, NodeReranker, SearchConfig +from graphiti_core.search.search_config_recipes import ( + EDGE_HYBRID_SEARCH_CROSS_ENCODER, + EDGE_HYBRID_SEARCH_NODE_DISTANCE, + EDGE_HYBRID_SEARCH_RRF, + NODE_HYBRID_SEARCH_CROSS_ENCODER, + NODE_HYBRID_SEARCH_NODE_DISTANCE, + NODE_HYBRID_SEARCH_RRF, +) +from graphiti_core.search.search_filters import ( + ComparisonOperator, + DateFilter, + SearchFilters, +) + +VALID_RERANKERS = {'rrf', 'mmr', 'cross_encoder'} + +# graphiti_core's maximal_marginal_relevance uses -2.0 as its "keep everything" +# sentinel. MMR scores span negatives (the diversity penalty can exceed the +# small query-similarity term), so any min_score >= 0 silently drops most +# candidates. reranker_min_score is only meaningful for the cross-encoder's +# calibrated 0-1 scores; for MMR we always disable filtering and let it reorder. +MMR_NO_FILTER_SCORE = -2.0 + + +def build_edge_search_config( + reranker: str, + mmr_lambda: float, + limit: int, + min_score: float, + center_node_uuid: str | None = None, +) -> SearchConfig: + if reranker not in VALID_RERANKERS: + raise ValueError(f'unknown reranker: {reranker!r}') + + if center_node_uuid is not None: + config = EDGE_HYBRID_SEARCH_NODE_DISTANCE.model_copy(deep=True) + elif reranker == 'cross_encoder': + config = EDGE_HYBRID_SEARCH_CROSS_ENCODER.model_copy(deep=True) + elif reranker == 'mmr': + config = EDGE_HYBRID_SEARCH_RRF.model_copy(deep=True) + assert config.edge_config is not None # RRF recipe always sets edge_config + config.edge_config.reranker = EdgeReranker.mmr + config.edge_config.mmr_lambda = mmr_lambda + min_score = MMR_NO_FILTER_SCORE # MMR must not filter; it only reorders + else: # 'rrf' + config = EDGE_HYBRID_SEARCH_RRF.model_copy(deep=True) + + config.limit = limit + config.reranker_min_score = min_score + return config + + +def build_node_search_config( + reranker: str, + mmr_lambda: float, + limit: int, + min_score: float, + center_node_uuid: str | None = None, +) -> SearchConfig: + if reranker not in VALID_RERANKERS: + raise ValueError(f'unknown reranker: {reranker!r}') + + if center_node_uuid is not None: + config = NODE_HYBRID_SEARCH_NODE_DISTANCE.model_copy(deep=True) + elif reranker == 'cross_encoder': + config = NODE_HYBRID_SEARCH_CROSS_ENCODER.model_copy(deep=True) + elif reranker == 'mmr': + config = NODE_HYBRID_SEARCH_RRF.model_copy(deep=True) + assert config.node_config is not None # RRF recipe always sets node_config + config.node_config.reranker = NodeReranker.mmr + config.node_config.mmr_lambda = mmr_lambda + min_score = MMR_NO_FILTER_SCORE # MMR must not filter; it only reorders + else: # 'rrf' + config = NODE_HYBRID_SEARCH_RRF.model_copy(deep=True) + + config.limit = limit + config.reranker_min_score = min_score + return config + + +def apply_liveness_filter(search_filter: SearchFilters | None) -> SearchFilters: + """Return a copy of ``search_filter`` that excludes superseded/expired edges. + + Adds ``expired_at IS NULL`` and ``(invalid_at IS NULL OR invalid_at > now)`` + unless the caller already constrained those fields (explicit caller intent + wins). Node search is unaffected — only edges carry temporal invalidation. + """ + sf = search_filter.model_copy(deep=True) if search_filter is not None else SearchFilters() + now = datetime.now(timezone.utc) + + if sf.expired_at is None: + sf.expired_at = [[DateFilter(comparison_operator=ComparisonOperator.is_null)]] + + if sf.invalid_at is None: + sf.invalid_at = [ + [DateFilter(comparison_operator=ComparisonOperator.is_null)], + [DateFilter(date=now, comparison_operator=ComparisonOperator.greater_than)], + ] + + return sf diff --git a/mcp_server/tests/test_formatting_score.py b/mcp_server/tests/test_formatting_score.py new file mode 100644 index 0000000000..18015f13ff --- /dev/null +++ b/mcp_server/tests/test_formatting_score.py @@ -0,0 +1,48 @@ +import sys +from datetime import datetime, timezone +from pathlib import Path + +sys.path.insert(0, str(Path(__file__).parent.parent / 'src')) + +from graphiti_core.edges import EntityEdge +from graphiti_core.nodes import EntityNode + +from utils.formatting import format_fact_result, to_node_result + + +def _node(): + return EntityNode( + uuid='n1', + name='FanWeb', + group_id='g', + labels=['Entity'], + created_at=datetime.now(timezone.utc), + summary='s', + ) + + +def _edge(): + return EntityEdge( + uuid='e1', + name='USES', + fact='A uses B', + group_id='g', + source_node_uuid='n1', + target_node_uuid='n2', + created_at=datetime.now(timezone.utc), + ) + + +def test_node_result_includes_score(): + r = to_node_result(_node(), score=0.42) + assert r['score'] == 0.42 + + +def test_node_result_score_defaults_none(): + r = to_node_result(_node()) + assert r['score'] is None + + +def test_fact_result_includes_score(): + r = format_fact_result(_edge(), score=0.99) + assert r['score'] == 0.99 diff --git a/mcp_server/tests/test_search_config_schema.py b/mcp_server/tests/test_search_config_schema.py new file mode 100644 index 0000000000..4505126836 --- /dev/null +++ b/mcp_server/tests/test_search_config_schema.py @@ -0,0 +1,52 @@ +import os +import sys +from pathlib import Path + +sys.path.insert(0, str(Path(__file__).parent.parent / 'src')) + +from config.schema import GraphitiConfig, SearchTuningConfig + + +def test_search_defaults(): + cfg = GraphitiConfig() + assert isinstance(cfg.search, SearchTuningConfig) + assert cfg.search.reranker == 'rrf' + assert cfg.search.mmr_lambda == 0.5 + assert cfg.search.max_facts == 6 + assert cfg.search.max_nodes == 6 + assert cfg.search.exclude_invalidated is True + assert cfg.search.reranker_min_score == 0.0 + + +def test_search_env_overrides(): + os.environ['SEARCH__RERANKER'] = 'mmr' + os.environ['SEARCH__MMR_LAMBDA'] = '0.7' + os.environ['SEARCH__MAX_FACTS'] = '4' + os.environ['SEARCH__EXCLUDE_INVALIDATED'] = 'false' + try: + cfg = GraphitiConfig() + assert cfg.search.reranker == 'mmr' + assert cfg.search.mmr_lambda == 0.7 + assert cfg.search.max_facts == 4 + assert cfg.search.exclude_invalidated is False + finally: + for k in ( + 'SEARCH__RERANKER', + 'SEARCH__MMR_LAMBDA', + 'SEARCH__MAX_FACTS', + 'SEARCH__EXCLUDE_INVALIDATED', + ): + os.environ.pop(k, None) + + +def test_invalid_reranker_rejected(): + os.environ['SEARCH__RERANKER'] = 'bogus' + try: + raised = False + try: + GraphitiConfig() + except Exception: + raised = True + assert raised, 'expected validation error for unknown reranker' + finally: + os.environ.pop('SEARCH__RERANKER', None) diff --git a/mcp_server/tests/test_search_relevance_int.py b/mcp_server/tests/test_search_relevance_int.py new file mode 100644 index 0000000000..b7c701f25b --- /dev/null +++ b/mcp_server/tests/test_search_relevance_int.py @@ -0,0 +1,138 @@ +"""Live-graph relevance checks for the tuned MCP search tools (opt-in). + +Proves the behavioral claims that unit tests cannot: against a real, populated +graph, (a) no superseded/expired facts are returned by default, and (b) every +returned fact/node carries a rerank ``score``. + +Opt-in: set ``GRAPHITI_INT_GROUP_ID`` to the group to query. The server is +spawned as a subprocess over stdio (same bootstrap as +``test_live_falkordb_int.py``), so it also needs a reachable FalkorDB and a +working embedder/LLM config (matching whatever populated the target graph). + +Example:: + + cd mcp_server + GRAPHITI_INT_GROUP_ID=brent_atvenu uv run pytest tests/test_search_relevance_int.py -v +""" + +import json +import os +import socket +from contextlib import AsyncExitStack, suppress +from datetime import datetime, timezone +from pathlib import Path +from typing import Any + +import pytest +from dotenv import load_dotenv +from mcp import ClientSession, StdioServerParameters +from mcp.client.stdio import stdio_client + +_REPO_ROOT = Path(__file__).resolve().parents[2] +load_dotenv(_REPO_ROOT / '.env') + +MCP_SERVER_DIR = Path(__file__).resolve().parent.parent +FALKORDB_URI = os.environ.get('FALKORDB_URI', 'redis://localhost:6379') + +pytestmark = [ + pytest.mark.integration, + pytest.mark.requires_falkordb, +] + + +def _falkordb_reachable(uri: str) -> bool: + host, port = 'localhost', 6379 + if '://' in uri: + rest = uri.split('://', 1)[1].split('@')[-1] + hostport = rest.split('/')[0] + if ':' in hostport: + host_part, port_part = hostport.rsplit(':', 1) + host = host_part or host + with suppress(ValueError): + port = int(port_part) + elif hostport: + host = hostport + try: + with socket.create_connection((host, port), timeout=2): + return True + except OSError: + return False + + +_GROUP = os.environ.get('GRAPHITI_INT_GROUP_ID') +if not _GROUP: + pytest.skip( + 'set GRAPHITI_INT_GROUP_ID to run live-graph relevance checks', + allow_module_level=True, + ) +if not _falkordb_reachable(FALKORDB_URI): + pytest.skip( + f'FalkorDB not reachable at {FALKORDB_URI}; skipping live relevance checks', + allow_module_level=True, + ) + + +class _Client: + """Minimal stdio MCP client for the running server.""" + + def __init__(self) -> None: + self._stack = AsyncExitStack() + self.session: ClientSession | None = None + + async def __aenter__(self) -> '_Client': + env = {**os.environ, 'FALKORDB_URI': FALKORDB_URI} + params = StdioServerParameters( + command='uv', + args=['run', str(MCP_SERVER_DIR / 'main.py'), '--transport', 'stdio'], + env=env, + cwd=str(MCP_SERVER_DIR), + ) + read, write = await self._stack.enter_async_context(stdio_client(params)) + self.session = await self._stack.enter_async_context(ClientSession(read, write)) + await self.session.initialize() + return self + + async def __aexit__(self, *exc: Any) -> None: + await self._stack.aclose() + + async def call(self, tool: str, arguments: dict[str, Any]) -> Any: + assert self.session is not None + result = await self.session.call_tool(tool, arguments) + if not result.content: + return None + text = getattr(result.content[0], 'text', None) + if text is None: + return None + try: + return json.loads(text) + except (json.JSONDecodeError, TypeError): + return text + + +async def test_no_expired_edges_returned_int(): + """Every returned fact must be live (invalid_at null or in the future) and scored.""" + async with _Client() as client: + resp = await client.call( + 'search_memory_facts', + {'query': 'Rails 6.1 upgrade', 'group_ids': [_GROUP], 'max_facts': 10}, + ) + assert isinstance(resp, dict) and not resp.get('error'), f'unexpected response: {resp}' + now = datetime.now(timezone.utc) + facts = resp.get('facts', []) + for fact in facts: + assert 'score' in fact, f'fact missing score: {fact.get("uuid")}' + inv = fact.get('invalid_at') + if inv: + assert datetime.fromisoformat(inv) > now, f'stale fact returned: {fact["uuid"]}' + + +async def test_nodes_carry_score_int(): + """Node search results carry a rerank score.""" + async with _Client() as client: + resp = await client.call( + 'search_nodes', + {'query': 'Rails upgrade', 'group_ids': [_GROUP], 'max_nodes': 10}, + ) + assert isinstance(resp, dict) and not resp.get('error'), f'unexpected response: {resp}' + for node in resp.get('nodes', []): + assert 'score' in node, f'node missing score: {node.get("uuid")}' diff --git a/mcp_server/tests/test_search_tools_wiring.py b/mcp_server/tests/test_search_tools_wiring.py new file mode 100644 index 0000000000..af893a4b05 --- /dev/null +++ b/mcp_server/tests/test_search_tools_wiring.py @@ -0,0 +1,100 @@ +import asyncio +import sys +from datetime import datetime, timezone +from pathlib import Path + +sys.path.insert(0, str(Path(__file__).parent.parent / 'src')) + +from graphiti_core.edges import EntityEdge +from graphiti_core.nodes import EntityNode +from graphiti_core.search.search_config import EdgeReranker, NodeReranker, SearchResults + +import graphiti_mcp_server as srv +from config.schema import GraphitiConfig + + +class _FakeClient: + def __init__(self): + self.calls = [] + + async def search_(self, *, query, config, group_ids, center_node_uuid, search_filter): + self.calls.append( + {'config': config, 'search_filter': search_filter, 'group_ids': group_ids} + ) + edge = EntityEdge( + uuid='e1', + name='USES', + fact='A uses B', + group_id='g', + source_node_uuid='n1', + target_node_uuid='n2', + created_at=datetime.now(timezone.utc), + ) + node = EntityNode( + uuid='n1', + name='FanWeb', + group_id='g', + labels=['Entity'], + created_at=datetime.now(timezone.utc), + summary='s', + ) + return SearchResults( + edges=[edge], + nodes=[node], + edge_reranker_scores=[0.87], + node_reranker_scores=[0.73], + ) + + +class _FakeService: + def __init__(self, client): + self._client = client + + async def get_client(self): + return self._client + + +def _install_fake(monkeypatch): + client = _FakeClient() + monkeypatch.setattr(srv, 'graphiti_service', _FakeService(client)) + # The module-level ``config`` global is only assigned inside initialize_server(); + # the tool functions read config.search / config.graphiti, so provide a default. + monkeypatch.setattr(srv, 'config', GraphitiConfig(), raising=False) + return client + + +def test_facts_use_default_reranker_and_surface_score(monkeypatch): + client = _install_fake(monkeypatch) + resp = asyncio.run(srv.search_memory_facts(query='q', group_ids='g')) + assert resp['facts'][0]['score'] == 0.87 + call = client.calls[0] + # rrf is the default reranker (see MMR finding: mmr is miscalibrated here) + assert call['config'].edge_config.reranker == EdgeReranker.rrf + # liveness filter applied by default + assert call['search_filter'].expired_at is not None + + +def test_facts_include_invalidated_opt_in(monkeypatch): + client = _install_fake(monkeypatch) + asyncio.run(srv.search_memory_facts(query='q', group_ids='g', include_invalidated=True)) + call = client.calls[0] + assert call['search_filter'].expired_at is None + + +def test_facts_default_limit_is_six(monkeypatch): + client = _install_fake(monkeypatch) + asyncio.run(srv.search_memory_facts(query='q', group_ids='g')) + assert client.calls[0]['config'].limit == 6 + + +def test_nodes_use_default_reranker_and_surface_score(monkeypatch): + client = _install_fake(monkeypatch) + resp = asyncio.run(srv.search_nodes(query='q', group_ids='g')) + assert resp['nodes'][0]['score'] == 0.73 + assert client.calls[0]['config'].node_config.reranker == NodeReranker.rrf + + +def test_nodes_default_limit_is_six(monkeypatch): + client = _install_fake(monkeypatch) + asyncio.run(srv.search_nodes(query='q', group_ids='g')) + assert client.calls[0]['config'].limit == 6 diff --git a/mcp_server/tests/test_search_tuning.py b/mcp_server/tests/test_search_tuning.py new file mode 100644 index 0000000000..421dfb6925 --- /dev/null +++ b/mcp_server/tests/test_search_tuning.py @@ -0,0 +1,98 @@ +import sys +from pathlib import Path + +sys.path.insert(0, str(Path(__file__).parent.parent / 'src')) + +from graphiti_core.search.search_config import EdgeReranker, NodeReranker +from graphiti_core.search.search_config_recipes import EDGE_HYBRID_SEARCH_RRF +from graphiti_core.search.search_filters import ComparisonOperator, DateFilter, SearchFilters + +from utils.search_tuning import ( + apply_liveness_filter, + build_edge_search_config, + build_node_search_config, +) + + +def test_edge_mmr_sets_lambda_and_reranker(): + cfg = build_edge_search_config('mmr', mmr_lambda=0.5, limit=6, min_score=0.0) + assert cfg.edge_config.reranker == EdgeReranker.mmr + assert cfg.edge_config.mmr_lambda == 0.5 + assert cfg.limit == 6 + # MMR must never filter: the caller's min_score (0.0 here) is overridden with + # the no-filter sentinel, or MMR's mostly-negative scores would drop results. + assert cfg.reranker_min_score == -2.0 + + +def test_node_mmr_forces_no_filter_min_score(): + cfg = build_node_search_config('mmr', mmr_lambda=0.5, limit=6, min_score=0.0) + assert cfg.reranker_min_score == -2.0 + + +def test_edge_rrf_keeps_rrf(): + cfg = build_edge_search_config('rrf', mmr_lambda=0.5, limit=4, min_score=0.0) + assert cfg.edge_config.reranker == EdgeReranker.rrf + assert cfg.limit == 4 + + +def test_edge_cross_encoder_sets_min_score(): + cfg = build_edge_search_config('cross_encoder', mmr_lambda=0.5, limit=6, min_score=0.2) + assert cfg.edge_config.reranker == EdgeReranker.cross_encoder + assert cfg.reranker_min_score == 0.2 + + +def test_edge_center_node_forces_node_distance(): + cfg = build_edge_search_config( + 'mmr', mmr_lambda=0.5, limit=6, min_score=0.0, center_node_uuid='abc' + ) + assert cfg.edge_config.reranker == EdgeReranker.node_distance + + +def test_node_mmr(): + cfg = build_node_search_config('mmr', mmr_lambda=0.5, limit=6, min_score=0.0) + assert cfg.node_config.reranker == NodeReranker.mmr + assert cfg.node_config.mmr_lambda == 0.5 + + +def test_does_not_mutate_singleton(): + before = EDGE_HYBRID_SEARCH_RRF.limit + _ = build_edge_search_config('mmr', mmr_lambda=0.3, limit=999, min_score=0.0) + assert EDGE_HYBRID_SEARCH_RRF.limit == before + assert EDGE_HYBRID_SEARCH_RRF.edge_config.reranker == EdgeReranker.rrf + + +def test_unknown_reranker_raises(): + raised = False + try: + build_edge_search_config('bogus', mmr_lambda=0.5, limit=6, min_score=0.0) + except ValueError: + raised = True + assert raised + + +def test_liveness_filter_from_none(): + sf = apply_liveness_filter(None) + # expired_at: single AND-group with one IS NULL clause + assert sf.expired_at is not None + assert sf.expired_at[0][0].comparison_operator == ComparisonOperator.is_null + # invalid_at: two OR-groups (IS NULL) OR (> now) + assert sf.invalid_at is not None + assert len(sf.invalid_at) == 2 + assert sf.invalid_at[0][0].comparison_operator == ComparisonOperator.is_null + assert sf.invalid_at[1][0].comparison_operator == ComparisonOperator.greater_than + + +def test_liveness_filter_preserves_explicit_expired_at(): + explicit = SearchFilters( + expired_at=[[DateFilter(comparison_operator=ComparisonOperator.is_not_null)]] + ) + sf = apply_liveness_filter(explicit) + # caller intent wins: we do NOT overwrite an explicit expired_at + assert sf.expired_at[0][0].comparison_operator == ComparisonOperator.is_not_null + + +def test_liveness_filter_preserves_edge_types(): + explicit = SearchFilters(edge_types=['FIXED_BY']) + sf = apply_liveness_filter(explicit) + assert sf.edge_types == ['FIXED_BY'] + assert sf.expired_at is not None # liveness still added