mirror of
https://github.com/basicmachines-co/basic-memory
synced 2026-06-21 13:47:35 +00:00
fix tests, add SearchResult response type for search
This commit is contained in:
phernandez
@@ -24,7 +24,7 @@ from basic_memory.mcp.tools.knowledge import (
|
||||
)
|
||||
|
||||
from basic_memory.mcp.tools.search import (
|
||||
search_nodes,
|
||||
search,
|
||||
open_nodes,
|
||||
)
|
||||
|
||||
@@ -56,7 +56,7 @@ __all__ = [
|
||||
"delete_relations",
|
||||
|
||||
# Search tools
|
||||
"search_nodes",
|
||||
"search",
|
||||
"get_entity",
|
||||
"open_nodes",
|
||||
|
||||
|
||||
@@ -1,186 +1,151 @@
|
||||
"""Search and query tools for Basic Memory MCP server."""
|
||||
"""Search tools for Basic Memory MCP server."""
|
||||
|
||||
from typing import List, Optional
|
||||
from datetime import datetime, timezone
|
||||
import textwrap
|
||||
from collections import defaultdict
|
||||
|
||||
from basic_memory.mcp.server import mcp
|
||||
from basic_memory.schemas.request import SearchNodesRequest, OpenNodesRequest
|
||||
from basic_memory.schemas.response import SearchNodesResponse, EntityListResponse
|
||||
from basic_memory.schemas.search import SearchQuery, SearchResult, SearchItemType
|
||||
from basic_memory.schemas.request import OpenNodesRequest
|
||||
from basic_memory.schemas.response import EntityListResponse
|
||||
from basic_memory.mcp.async_client import client
|
||||
|
||||
|
||||
@mcp.tool(
|
||||
category="search",
|
||||
description="Search for entities across names, descriptions, observations, and relations",
|
||||
description="Search across all content in basic-memory, including documents and entities",
|
||||
examples=[
|
||||
{
|
||||
"name": "Technical Search",
|
||||
"description": "Find implementation details and patterns",
|
||||
"name": "Search with Metadata Analysis",
|
||||
"description": "Search and analyze results by metadata",
|
||||
"code": """
|
||||
# Search for database-related components
|
||||
results = await search_nodes(
|
||||
request=SearchNodesRequest(
|
||||
query="sqlite database implementation",
|
||||
category="tech" # Focus on technical details
|
||||
)
|
||||
# Search for feature specs
|
||||
results = await search(
|
||||
text="implementation",
|
||||
types=[SearchItemType.DOCUMENT]
|
||||
)
|
||||
|
||||
# Analyze implementation patterns
|
||||
for entity in results.matches:
|
||||
print(f"\\n{entity.name} Implementation:")
|
||||
|
||||
# Technical details
|
||||
tech_notes = [o.content for o in entity.observations
|
||||
if o.category == "tech"]
|
||||
if tech_notes:
|
||||
print("Technical Notes:")
|
||||
for note in tech_notes:
|
||||
print(f"- {note}")
|
||||
|
||||
# Dependencies
|
||||
deps = [r for r in entity.relations
|
||||
if r.relation_type == "depends_on"]
|
||||
if deps:
|
||||
print("\\nDependencies:")
|
||||
for dep in deps:
|
||||
print(f"- {dep.to_id}")"""
|
||||
# Group by category and status
|
||||
by_category = defaultdict(list)
|
||||
by_status = defaultdict(list)
|
||||
|
||||
for r in results:
|
||||
meta = r.metadata
|
||||
if 'category' in meta:
|
||||
by_category[meta['category']].append(r)
|
||||
if 'status' in meta:
|
||||
by_status[meta['status']].append(r)
|
||||
|
||||
print("Results by Category:")
|
||||
for category, items in by_category.items():
|
||||
print(f"\\n{category.title()}:")
|
||||
for item in items:
|
||||
print(f"- {item.path_id} (score: {item.score:.2f})")
|
||||
|
||||
# Find high priority items
|
||||
high_priority = [
|
||||
r for r in results
|
||||
if r.metadata.get('priority') in ['high', 'highest']
|
||||
]
|
||||
"""
|
||||
},
|
||||
{
|
||||
"name": "Feature Context",
|
||||
"description": "Build complete feature implementation context",
|
||||
"name": "Recent Changes Analysis",
|
||||
"description": "Search and analyze recent document changes",
|
||||
"code": """
|
||||
# Start with feature search
|
||||
feature_results = await search_nodes(
|
||||
request=SearchNodesRequest(
|
||||
query="semantic search feature"
|
||||
)
|
||||
from datetime import datetime, timedelta
|
||||
|
||||
# Set cutoff date
|
||||
cutoff = datetime.now(timezone.utc) - timedelta(days=7)
|
||||
|
||||
# Search for recent changes
|
||||
results = await search(text="database", after_date=cutoff)
|
||||
|
||||
# Sort by update time
|
||||
sorted_results = sorted(
|
||||
results,
|
||||
key=lambda x: x.metadata['updated_at'],
|
||||
reverse=True
|
||||
)
|
||||
|
||||
# Collect related entities for context
|
||||
related_ids = set()
|
||||
for entity in feature_results.matches:
|
||||
# Add feature itself
|
||||
related_ids.add(entity.path_id)
|
||||
# Add related entities
|
||||
for relation in entity.relations:
|
||||
related_ids.add(relation.to_id)
|
||||
print("Recent Changes:")
|
||||
for r in sorted_results[:5]:
|
||||
print(f"\\n{r.path_id}")
|
||||
print(f"Updated: {r.metadata['updated_at']}")
|
||||
if 'author' in r.metadata:
|
||||
print(f"Author: {r.metadata['author']}")
|
||||
print(f"Score: {r.score:.2f}")
|
||||
"""
|
||||
},
|
||||
{
|
||||
"name": "Entity Context Loading",
|
||||
"description": "Search for entities and load their full context",
|
||||
"code": """
|
||||
# Find relevant components
|
||||
results = await search(
|
||||
text="knowledge graph",
|
||||
types=[SearchItemType.ENTITY],
|
||||
entity_types=["component"]
|
||||
)
|
||||
|
||||
# Load complete context
|
||||
if related_ids:
|
||||
if results:
|
||||
# Load full entity details
|
||||
path_ids = [r.path_id for r in results]
|
||||
context = await open_nodes(
|
||||
request=OpenNodesRequest(
|
||||
path_ids=list(related_ids)
|
||||
)
|
||||
request=OpenNodesRequest(path_ids=path_ids)
|
||||
)
|
||||
|
||||
# Analyze implementation status
|
||||
components = [e for e in context.entities
|
||||
if e.entity_type == "component"]
|
||||
tests = [e for e in context.entities
|
||||
if e.entity_type == "test"]
|
||||
specs = [e for e in context.entities
|
||||
if e.entity_type == "specification"]
|
||||
|
||||
print("Implementation Status:")
|
||||
print(f"- Components: {len(components)}")
|
||||
print(f"- Tests: {len(tests)}")
|
||||
print(f"- Specs: {len(specs)}")"""
|
||||
},
|
||||
{
|
||||
"name": "Design Analysis",
|
||||
"description": "Extract architectural decisions and patterns",
|
||||
"code": """
|
||||
# Search for design decisions
|
||||
design_results = await search_nodes(
|
||||
request=SearchNodesRequest(
|
||||
query="architecture pattern",
|
||||
category="design"
|
||||
)
|
||||
)
|
||||
|
||||
# Group decisions by component
|
||||
from collections import defaultdict
|
||||
decisions = defaultdict(list)
|
||||
|
||||
for entity in design_results.matches:
|
||||
# Extract design observations
|
||||
design_notes = [o for o in entity.observations
|
||||
if o.category == "design"]
|
||||
if design_notes:
|
||||
decisions[entity.name].extend(design_notes)
|
||||
|
||||
# Show architectural decisions
|
||||
for component, notes in decisions.items():
|
||||
print(f"\\n{component} Architecture:")
|
||||
for note in notes:
|
||||
context = note.context or "Design Decision"
|
||||
print(f"\\n{context}:")
|
||||
print(f"- {note.content}")"""
|
||||
},
|
||||
{
|
||||
"name": "Knowledge Chain",
|
||||
"description": "Follow knowledge links to build deep context",
|
||||
"code": """
|
||||
# Start with initial concept
|
||||
initial = await search_nodes(
|
||||
request=SearchNodesRequest(query="semantic web")
|
||||
)
|
||||
|
||||
# Build knowledge chain
|
||||
seen_ids = set()
|
||||
to_explore = set()
|
||||
|
||||
# Add initial matches
|
||||
for entity in initial.matches:
|
||||
seen_ids.add(entity.path_id)
|
||||
for relation in entity.relations:
|
||||
to_explore.add(relation.to_id)
|
||||
|
||||
# Explore up to 2 levels deep
|
||||
knowledge_chain = initial.matches
|
||||
for _ in range(2):
|
||||
if not to_explore:
|
||||
break
|
||||
# Analyze implementation details
|
||||
print("Implementation Components:")
|
||||
for entity in context.entities:
|
||||
print(f"\\n{entity.name}")
|
||||
|
||||
# Load next level
|
||||
next_ids = list(to_explore - seen_ids)
|
||||
if next_ids:
|
||||
next_level = await open_nodes(
|
||||
request=OpenNodesRequest(path_ids=next_ids)
|
||||
)
|
||||
|
||||
# Update tracking
|
||||
knowledge_chain.extend(next_level.entities)
|
||||
seen_ids.update(next_ids)
|
||||
to_explore.clear()
|
||||
|
||||
# Add new relations
|
||||
for entity in next_level.entities:
|
||||
for relation in entity.relations:
|
||||
to_explore.add(relation.to_id)
|
||||
|
||||
# Analyze knowledge structure
|
||||
print(f"Knowledge chain depth: {len(seen_ids)} entities")
|
||||
type_counts = defaultdict(int)
|
||||
for entity in knowledge_chain:
|
||||
type_counts[entity.entity_type] += 1
|
||||
|
||||
print("\\nKnowledge composition:")
|
||||
for type_, count in type_counts.items():
|
||||
print(f"- {type_}: {count} entities")"""
|
||||
# Show technical details
|
||||
tech_notes = [
|
||||
o.content for o in entity.observations
|
||||
if o.category == 'tech'
|
||||
]
|
||||
if tech_notes:
|
||||
print("Technical Notes:")
|
||||
for note in tech_notes:
|
||||
print(f"- {note}")
|
||||
|
||||
# Show dependencies
|
||||
deps = [r for r in entity.relations if r.relation_type == 'depends_on']
|
||||
if deps:
|
||||
print("\\nDependencies:")
|
||||
for dep in deps:
|
||||
print(f"- {dep.to_id}")
|
||||
"""
|
||||
}
|
||||
],
|
||||
output_model=SearchNodesResponse,
|
||||
]
|
||||
)
|
||||
async def search_nodes(request: SearchNodesRequest) -> SearchNodesResponse:
|
||||
"""Search for entities in the knowledge graph.
|
||||
async def search(
|
||||
text: str,
|
||||
types: Optional[List[SearchItemType]] = None,
|
||||
entity_types: Optional[List[str]] = None,
|
||||
after_date: Optional[datetime] = None
|
||||
) -> List[SearchResult]:
|
||||
"""Search across all content in basic-memory.
|
||||
|
||||
Args:
|
||||
request: Search parameters including query text and optional category
|
||||
text: Text to search for
|
||||
types: Optional list of types to filter by (DOCUMENT, ENTITY)
|
||||
entity_types: Optional list of entity types to filter by
|
||||
after_date: Optional date to filter results after
|
||||
|
||||
Returns:
|
||||
SearchNodesResponse containing matching entities and search metadata
|
||||
List of SearchResult objects sorted by relevance
|
||||
"""
|
||||
url = "/knowledge/search"
|
||||
response = await client.post(url, json=request.model_dump())
|
||||
return SearchNodesResponse.model_validate(response.json())
|
||||
query = SearchQuery(
|
||||
text=text,
|
||||
types=types,
|
||||
entity_types=entity_types,
|
||||
after_date=after_date
|
||||
)
|
||||
response = await client.post("/search/", json=query.model_dump())
|
||||
return [SearchResult.model_validate(r) for r in response.json()]
|
||||
|
||||
|
||||
@mcp.tool(
|
||||
@@ -188,100 +153,39 @@ async def search_nodes(request: SearchNodesRequest) -> SearchNodesResponse:
|
||||
description="Load multiple entities by their path_ids in a single request",
|
||||
examples=[
|
||||
{
|
||||
"name": "Implementation Chain",
|
||||
"description": "Load and analyze implementation dependencies",
|
||||
"name": "Load Search Context",
|
||||
"description": "Load full entity details from search results",
|
||||
"code": """
|
||||
# Load feature implementation chain
|
||||
chain = await open_nodes(
|
||||
request=OpenNodesRequest(
|
||||
path_ids=[
|
||||
"feature/semantic_search", # The feature
|
||||
"component/search_service", # Core implementation
|
||||
"component/index_service", # Supporting service
|
||||
"test/search_integration", # Integration tests
|
||||
"document/search_spec" # Documentation
|
||||
]
|
||||
)
|
||||
# First search for entities
|
||||
results = await search(
|
||||
text="database implementation",
|
||||
types=[SearchItemType.ENTITY]
|
||||
)
|
||||
|
||||
def analyze_dependencies(entities):
|
||||
deps = defaultdict(list)
|
||||
for entity in entities:
|
||||
# Direct dependencies
|
||||
direct = [r.to_id for r in entity.relations
|
||||
if r.relation_type == "depends_on"]
|
||||
deps[entity.path_id].extend(direct)
|
||||
|
||||
# Implicit dependencies via observations
|
||||
for obs in entity.observations:
|
||||
if "requires" in obs.content.lower():
|
||||
deps[entity.path_id].append(
|
||||
f"Implicit: {obs.content}"
|
||||
)
|
||||
return deps
|
||||
|
||||
# Show implementation structure
|
||||
deps = analyze_dependencies(chain.entities)
|
||||
for path_id, dependencies in deps.items():
|
||||
print(f"\\n{path_id} dependencies:")
|
||||
for dep in dependencies:
|
||||
print(f"- {dep}")"""
|
||||
},
|
||||
{
|
||||
"name": "Technical Analysis",
|
||||
"description": "Deep dive into technical implementation",
|
||||
"code": """
|
||||
# First find technical components
|
||||
tech_results = await search_nodes(
|
||||
request=SearchNodesRequest(
|
||||
query="search implementation",
|
||||
category="tech"
|
||||
)
|
||||
)
|
||||
|
||||
# Load full technical context
|
||||
tech_ids = [e.path_id for e in tech_results.matches
|
||||
if e.entity_type == "component"]
|
||||
|
||||
if tech_ids:
|
||||
details = await open_nodes(
|
||||
request=OpenNodesRequest(path_ids=tech_ids)
|
||||
# Then load full context
|
||||
if results:
|
||||
path_ids = [r.path_id for r in results]
|
||||
context = await open_nodes(
|
||||
request=OpenNodesRequest(path_ids=path_ids)
|
||||
)
|
||||
|
||||
# Analyze technical architecture
|
||||
print("Technical Architecture:\\n")
|
||||
|
||||
for entity in details.entities:
|
||||
print(f"{entity.name}:")
|
||||
# Group by entity type
|
||||
by_type = defaultdict(list)
|
||||
for entity in context.entities:
|
||||
by_type[entity.entity_type].append(entity)
|
||||
|
||||
# Core capabilities
|
||||
tech_notes = [o.content for o in entity.observations
|
||||
if o.category == "tech"]
|
||||
if tech_notes:
|
||||
print("\\nCapabilities:")
|
||||
for note in tech_notes:
|
||||
print(f"- {note}")
|
||||
|
||||
# Design decisions
|
||||
design_notes = [o.content for o in entity.observations
|
||||
if o.category == "design"]
|
||||
if design_notes:
|
||||
print("\\nDesign Decisions:")
|
||||
for note in design_notes:
|
||||
print(f"- {note}")
|
||||
|
||||
# Dependencies
|
||||
deps = [r for r in entity.relations
|
||||
if r.relation_type == "depends_on"]
|
||||
if deps:
|
||||
print("\\nDependencies:")
|
||||
for dep in deps:
|
||||
print(f"- {dep.to_id}")
|
||||
|
||||
print("\\n---")"""
|
||||
# Show breakdown
|
||||
for etype, entities in by_type.items():
|
||||
print(f"\\n{etype.title()} Components:")
|
||||
for entity in entities:
|
||||
print(f"- {entity.name}")
|
||||
if entity.observations:
|
||||
print(f" {len(entity.observations)} observations")
|
||||
if entity.relations:
|
||||
print(f" {len(entity.relations)} relations")
|
||||
"""
|
||||
}
|
||||
],
|
||||
output_model=EntityListResponse,
|
||||
]
|
||||
)
|
||||
async def open_nodes(request: OpenNodesRequest) -> EntityListResponse:
|
||||
"""Load multiple entities by their path_ids.
|
||||
|
||||
Reference in New Issue
Block a user