mirror of
https://github.com/basicmachines-co/basic-memory
synced 2026-06-21 13:47:35 +00:00
chore(core): use ty for typechecking
Signed-off-by: phernandez <paul@basicmachines.co>
This commit is contained in:
@@ -6,6 +6,7 @@ have entity IDs in URLs - they generate formatted prompts from queries.
|
||||
"""
|
||||
|
||||
from datetime import datetime, timezone
|
||||
from typing import Any
|
||||
from fastapi import APIRouter, HTTPException, status, Path
|
||||
from loguru import logger
|
||||
|
||||
@@ -59,6 +60,7 @@ async def continue_conversation(
|
||||
|
||||
# Initialize search results
|
||||
search_results = []
|
||||
hierarchical_results_for_count = []
|
||||
|
||||
# Get data needed for template
|
||||
if request.topic:
|
||||
@@ -91,7 +93,8 @@ async def continue_conversation(
|
||||
# Limit to a reasonable number of total results
|
||||
all_hierarchical_results = all_hierarchical_results[:10]
|
||||
|
||||
template_context = {
|
||||
hierarchical_results_for_count = all_hierarchical_results
|
||||
template_context: dict[str, Any] = {
|
||||
"topic": request.topic,
|
||||
"timeframe": request.timeframe,
|
||||
"hierarchical_results": all_hierarchical_results,
|
||||
@@ -110,6 +113,7 @@ async def continue_conversation(
|
||||
|
||||
hierarchical_results = recent_context.results[:5] # Limit to top 5 recent items
|
||||
|
||||
hierarchical_results_for_count = hierarchical_results
|
||||
template_context = {
|
||||
"topic": f"Recent Activity from ({request.timeframe})",
|
||||
"timeframe": request.timeframe,
|
||||
@@ -129,9 +133,6 @@ async def continue_conversation(
|
||||
relation_count = 0
|
||||
entity_count = 0
|
||||
|
||||
# Get the hierarchical results from the template context
|
||||
hierarchical_results_for_count = template_context.get("hierarchical_results", [])
|
||||
|
||||
# For topic-based search
|
||||
if request.topic:
|
||||
for item in hierarchical_results_for_count:
|
||||
@@ -159,29 +160,24 @@ async def continue_conversation(
|
||||
elif related.type == "entity": # pragma: no cover
|
||||
entity_count += 1 # pragma: no cover
|
||||
|
||||
# Build metadata
|
||||
metadata = {
|
||||
"query": request.topic,
|
||||
"timeframe": request.timeframe,
|
||||
"search_count": len(search_results)
|
||||
if request.topic
|
||||
else 0, # Original search results count
|
||||
"context_count": len(hierarchical_results_for_count),
|
||||
"observation_count": observation_count,
|
||||
"relation_count": relation_count,
|
||||
"total_items": (
|
||||
prompt_metadata = PromptMetadata(
|
||||
query=request.topic,
|
||||
timeframe=request.timeframe,
|
||||
search_count=len(search_results) if request.topic else 0,
|
||||
context_count=len(hierarchical_results_for_count),
|
||||
observation_count=observation_count,
|
||||
relation_count=relation_count,
|
||||
total_items=(
|
||||
len(hierarchical_results_for_count)
|
||||
+ observation_count
|
||||
+ relation_count
|
||||
+ entity_count
|
||||
),
|
||||
"search_limit": request.search_items_limit,
|
||||
"context_depth": request.depth,
|
||||
"related_limit": request.related_items_limit,
|
||||
"generated_at": datetime.now(timezone.utc).isoformat(),
|
||||
}
|
||||
|
||||
prompt_metadata = PromptMetadata(**metadata)
|
||||
search_limit=request.search_items_limit,
|
||||
context_depth=request.depth,
|
||||
related_limit=request.related_items_limit,
|
||||
generated_at=datetime.now(timezone.utc).isoformat(),
|
||||
)
|
||||
|
||||
return PromptResponse(
|
||||
prompt=rendered_prompt, context=template_context, metadata=prompt_metadata
|
||||
@@ -229,7 +225,7 @@ async def search_prompt(
|
||||
results = await search_service.search(query, limit=limit, offset=offset)
|
||||
search_results = await to_search_results(entity_service, results)
|
||||
|
||||
template_context = {
|
||||
template_context: dict[str, Any] = {
|
||||
"query": request.query,
|
||||
"timeframe": request.timeframe,
|
||||
"results": search_results,
|
||||
@@ -241,22 +237,19 @@ async def search_prompt(
|
||||
# Render template
|
||||
rendered_prompt = await template_loader.render("prompts/search.hbs", template_context)
|
||||
|
||||
# Build metadata
|
||||
metadata = {
|
||||
"query": request.query,
|
||||
"timeframe": request.timeframe,
|
||||
"search_count": len(search_results),
|
||||
"context_count": len(search_results),
|
||||
"observation_count": 0, # Search results don't include observations
|
||||
"relation_count": 0, # Search results don't include relations
|
||||
"total_items": len(search_results),
|
||||
"search_limit": limit,
|
||||
"context_depth": 0, # No context depth for basic search
|
||||
"related_limit": 0, # No related items for basic search
|
||||
"generated_at": datetime.now(timezone.utc).isoformat(),
|
||||
}
|
||||
|
||||
prompt_metadata = PromptMetadata(**metadata)
|
||||
prompt_metadata = PromptMetadata(
|
||||
query=request.query,
|
||||
timeframe=request.timeframe,
|
||||
search_count=len(search_results),
|
||||
context_count=len(search_results),
|
||||
observation_count=0,
|
||||
relation_count=0,
|
||||
total_items=len(search_results),
|
||||
search_limit=limit,
|
||||
context_depth=0,
|
||||
related_limit=0,
|
||||
generated_at=datetime.now(timezone.utc).isoformat(),
|
||||
)
|
||||
|
||||
return PromptResponse(
|
||||
prompt=rendered_prompt, context=template_context, metadata=prompt_metadata
|
||||
|
||||
@@ -211,7 +211,7 @@ async def create_resource(
|
||||
action="create",
|
||||
phase="search_index",
|
||||
):
|
||||
await search_service.index_entity(entity) # pyright: ignore
|
||||
await search_service.index_entity(entity)
|
||||
|
||||
return ResourceResponse(
|
||||
entity_id=entity.id,
|
||||
@@ -326,6 +326,8 @@ async def update_resource(
|
||||
"updated_at": file_metadata.modified_at,
|
||||
},
|
||||
)
|
||||
if updated_entity is None:
|
||||
raise HTTPException(status_code=404, detail=f"Entity {entity_id} not found")
|
||||
|
||||
with telemetry.scope(
|
||||
"api.resource.update.search_index",
|
||||
@@ -333,7 +335,7 @@ async def update_resource(
|
||||
action="update",
|
||||
phase="search_index",
|
||||
):
|
||||
await search_service.index_entity(updated_entity) # pyright: ignore
|
||||
await search_service.index_entity(updated_entity)
|
||||
|
||||
return ResourceResponse(
|
||||
entity_id=entity.id,
|
||||
|
||||
@@ -1,8 +1,6 @@
|
||||
from typing import Optional, List
|
||||
from typing import Any, Protocol, Optional, List, Sequence
|
||||
|
||||
from basic_memory import telemetry
|
||||
from basic_memory.models import Entity as EntityModel
|
||||
from basic_memory.repository import EntityRepository
|
||||
from basic_memory.repository.search_repository import SearchIndexRow
|
||||
from basic_memory.schemas.memory import (
|
||||
EntitySummary,
|
||||
@@ -13,19 +11,38 @@ from basic_memory.schemas.memory import (
|
||||
ContextResult,
|
||||
)
|
||||
from basic_memory.schemas.search import SearchItemType, SearchResult
|
||||
from basic_memory.services import EntityService
|
||||
from basic_memory.services.context_service import (
|
||||
ContextResultRow,
|
||||
ContextResult as ServiceContextResult,
|
||||
)
|
||||
|
||||
|
||||
class EntityBatchLookup(Protocol):
|
||||
async def find_by_ids(self, ids: List[int]) -> Sequence[Any]: ...
|
||||
|
||||
|
||||
class EntityServiceBatchLookup(Protocol):
|
||||
async def get_entities_by_id(self, ids: List[int]) -> Sequence[Any]: ...
|
||||
|
||||
|
||||
def _required_str(value: str | None, field_name: str) -> str:
|
||||
"""Return a required search field or fail before producing invalid response data."""
|
||||
if value is None:
|
||||
raise ValueError(f"Search result is missing required field: {field_name}")
|
||||
return value
|
||||
|
||||
|
||||
def _search_item_type(value: str | SearchItemType) -> SearchItemType:
|
||||
"""Normalize repository row type strings into the public search enum."""
|
||||
return value if isinstance(value, SearchItemType) else SearchItemType(value)
|
||||
|
||||
|
||||
async def to_graph_context(
|
||||
context_result: ServiceContextResult,
|
||||
entity_repository: EntityRepository,
|
||||
entity_repository: EntityBatchLookup,
|
||||
page: Optional[int] = None,
|
||||
page_size: Optional[int] = None,
|
||||
):
|
||||
) -> GraphContext:
|
||||
with telemetry.scope(
|
||||
"memory.hydrate_context",
|
||||
domain="memory",
|
||||
@@ -44,17 +61,18 @@ async def to_graph_context(
|
||||
+ context_item.observations
|
||||
+ context_item.related_results
|
||||
):
|
||||
if item.type == SearchItemType.ENTITY:
|
||||
item_type = _search_item_type(item.type)
|
||||
if item_type == SearchItemType.ENTITY:
|
||||
# Entity's own ID for its external_id
|
||||
entity_ids_needed.add(item.id)
|
||||
elif item.type == SearchItemType.OBSERVATION:
|
||||
elif item_type == SearchItemType.OBSERVATION:
|
||||
# Parent entity ID for entity_external_id
|
||||
if item.entity_id: # pyright: ignore
|
||||
entity_ids_needed.add(item.entity_id) # pyright: ignore
|
||||
elif item.type == SearchItemType.RELATION:
|
||||
if item.entity_id:
|
||||
entity_ids_needed.add(item.entity_id)
|
||||
elif item_type == SearchItemType.RELATION:
|
||||
# Source and target entity IDs for external_ids
|
||||
if item.from_id: # pyright: ignore
|
||||
entity_ids_needed.add(item.from_id) # pyright: ignore
|
||||
if item.from_id:
|
||||
entity_ids_needed.add(item.from_id)
|
||||
if item.to_id:
|
||||
entity_ids_needed.add(item.to_id)
|
||||
|
||||
@@ -75,57 +93,60 @@ async def to_graph_context(
|
||||
entity_external_id_lookup[e.id] = e.external_id
|
||||
|
||||
# Helper function to convert items to summaries
|
||||
def to_summary(item: SearchIndexRow | ContextResultRow):
|
||||
match item.type:
|
||||
def to_summary(
|
||||
item: SearchIndexRow | ContextResultRow,
|
||||
) -> EntitySummary | ObservationSummary | RelationSummary:
|
||||
item_type = _search_item_type(item.type)
|
||||
match item_type:
|
||||
case SearchItemType.ENTITY:
|
||||
return EntitySummary(
|
||||
external_id=entity_external_id_lookup.get(item.id, ""),
|
||||
entity_id=item.id,
|
||||
title=item.title, # pyright: ignore
|
||||
title=_required_str(item.title, "title"),
|
||||
permalink=item.permalink,
|
||||
content=item.content,
|
||||
file_path=item.file_path,
|
||||
file_path=_required_str(item.file_path, "file_path"),
|
||||
created_at=item.created_at,
|
||||
)
|
||||
case SearchItemType.OBSERVATION:
|
||||
entity_ext_id = None
|
||||
if item.entity_id: # pyright: ignore
|
||||
entity_ext_id = entity_external_id_lookup.get(item.entity_id) # pyright: ignore
|
||||
entity_title = None
|
||||
if item.entity_id:
|
||||
entity_ext_id = entity_external_id_lookup.get(item.entity_id)
|
||||
entity_title = entity_title_lookup.get(item.entity_id)
|
||||
return ObservationSummary(
|
||||
observation_id=item.id,
|
||||
entity_id=item.entity_id, # pyright: ignore
|
||||
entity_id=item.entity_id,
|
||||
entity_external_id=entity_ext_id,
|
||||
title=entity_title_lookup.get(item.entity_id), # pyright: ignore
|
||||
file_path=item.file_path,
|
||||
category=item.category, # pyright: ignore
|
||||
content=item.content, # pyright: ignore
|
||||
permalink=item.permalink, # pyright: ignore
|
||||
title=entity_title,
|
||||
file_path=_required_str(item.file_path, "file_path"),
|
||||
category=_required_str(item.category, "category"),
|
||||
content=_required_str(item.content, "content"),
|
||||
permalink=_required_str(item.permalink, "permalink"),
|
||||
created_at=item.created_at,
|
||||
)
|
||||
case SearchItemType.RELATION:
|
||||
from_title = entity_title_lookup.get(item.from_id) if item.from_id else None # pyright: ignore
|
||||
from_title = entity_title_lookup.get(item.from_id) if item.from_id else None
|
||||
to_title = entity_title_lookup.get(item.to_id) if item.to_id else None
|
||||
from_ext_id = (
|
||||
entity_external_id_lookup.get(item.from_id) if item.from_id else None
|
||||
) # pyright: ignore
|
||||
)
|
||||
to_ext_id = entity_external_id_lookup.get(item.to_id) if item.to_id else None
|
||||
return RelationSummary(
|
||||
relation_id=item.id,
|
||||
entity_id=item.entity_id, # pyright: ignore
|
||||
title=item.title, # pyright: ignore
|
||||
file_path=item.file_path,
|
||||
permalink=item.permalink, # pyright: ignore
|
||||
relation_type=item.relation_type, # pyright: ignore
|
||||
entity_id=item.entity_id,
|
||||
title=_required_str(item.title, "title"),
|
||||
file_path=_required_str(item.file_path, "file_path"),
|
||||
permalink=_required_str(item.permalink, "permalink"),
|
||||
relation_type=_required_str(item.relation_type, "relation_type"),
|
||||
from_entity=from_title,
|
||||
from_entity_id=item.from_id, # pyright: ignore
|
||||
from_entity_id=item.from_id,
|
||||
from_entity_external_id=from_ext_id,
|
||||
to_entity=to_title,
|
||||
to_entity_id=item.to_id,
|
||||
to_entity_external_id=to_ext_id,
|
||||
created_at=item.created_at,
|
||||
)
|
||||
case _: # pragma: no cover
|
||||
raise ValueError(f"Unexpected type: {item.type}")
|
||||
|
||||
with telemetry.scope(
|
||||
"memory.hydrate_context.shape_results",
|
||||
@@ -137,12 +158,16 @@ async def to_graph_context(
|
||||
hierarchical_results = []
|
||||
for context_item in context_result.results:
|
||||
primary_result = to_summary(context_item.primary_result)
|
||||
observations = [to_summary(obs) for obs in context_item.observations]
|
||||
observations = [
|
||||
summary
|
||||
for summary in (to_summary(obs) for obs in context_item.observations)
|
||||
if isinstance(summary, ObservationSummary)
|
||||
]
|
||||
related = [to_summary(rel) for rel in context_item.related_results]
|
||||
hierarchical_results.append(
|
||||
ContextResult(
|
||||
primary_result=primary_result,
|
||||
observations=observations, # pyright: ignore[reportArgumentType]
|
||||
observations=observations,
|
||||
related_results=related,
|
||||
)
|
||||
)
|
||||
@@ -170,7 +195,9 @@ async def to_graph_context(
|
||||
)
|
||||
|
||||
|
||||
async def to_search_results(entity_service: EntityService, results: List[SearchIndexRow]):
|
||||
async def to_search_results(
|
||||
entity_service: EntityServiceBatchLookup, results: List[SearchIndexRow]
|
||||
) -> list[SearchResult]:
|
||||
with telemetry.scope(
|
||||
"search.hydrate_results",
|
||||
domain="search",
|
||||
@@ -187,7 +214,7 @@ async def to_search_results(entity_service: EntityService, results: List[SearchI
|
||||
all_entity_ids.add(eid)
|
||||
|
||||
# Single batch fetch for all entities
|
||||
entities_by_id: dict[int, EntityModel] = {}
|
||||
entities_by_id: dict[int, Any] = {}
|
||||
with telemetry.scope(
|
||||
"search.hydrate_results.fetch_entities",
|
||||
domain="search",
|
||||
@@ -222,20 +249,20 @@ async def to_search_results(entity_service: EntityService, results: List[SearchI
|
||||
entity_id = result.entity_id
|
||||
|
||||
# Look up entities by their specific IDs
|
||||
parent_entity = entities_by_id.get(result.entity_id) if result.entity_id else None # pyright: ignore
|
||||
from_entity = entities_by_id.get(result.from_id) if result.from_id else None # pyright: ignore
|
||||
parent_entity = entities_by_id.get(result.entity_id) if result.entity_id else None
|
||||
from_entity = entities_by_id.get(result.from_id) if result.from_id else None
|
||||
to_entity = entities_by_id.get(result.to_id) if result.to_id else None
|
||||
|
||||
search_results.append(
|
||||
SearchResult(
|
||||
title=result.title, # pyright: ignore
|
||||
type=result.type, # pyright: ignore
|
||||
title=_required_str(result.title, "title"),
|
||||
type=_search_item_type(result.type),
|
||||
permalink=result.permalink,
|
||||
score=result.score, # pyright: ignore
|
||||
score=result.score if result.score is not None else 0.0,
|
||||
entity=parent_entity.permalink if parent_entity else None,
|
||||
content=result.content,
|
||||
matched_chunk=result.matched_chunk_text,
|
||||
file_path=result.file_path,
|
||||
file_path=_required_str(result.file_path, "file_path"),
|
||||
metadata=result.metadata,
|
||||
entity_id=entity_id,
|
||||
observation_id=observation_id,
|
||||
|
||||
@@ -99,41 +99,39 @@ async def make_api_request(
|
||||
response = await client.request(method=method, url=url, headers=headers, json=json_data)
|
||||
response.raise_for_status()
|
||||
return response
|
||||
except httpx.HTTPStatusError as e:
|
||||
response = e.response
|
||||
|
||||
# Try to parse error detail from response
|
||||
error_detail = None
|
||||
try:
|
||||
error_detail = response.json()
|
||||
except Exception:
|
||||
# If JSON parsing fails, we'll handle it as a generic error
|
||||
pass
|
||||
|
||||
# Check for subscription_required error (403)
|
||||
if response.status_code == 403 and isinstance(error_detail, dict):
|
||||
# Handle both FastAPI HTTPException format (nested under "detail")
|
||||
# and direct format
|
||||
detail_obj = error_detail.get("detail", error_detail)
|
||||
if (
|
||||
isinstance(detail_obj, dict)
|
||||
and detail_obj.get("error") == "subscription_required"
|
||||
):
|
||||
message = detail_obj.get("message", "Active subscription required")
|
||||
subscribe_url = detail_obj.get(
|
||||
"subscribe_url", "https://basicmemory.com/subscribe"
|
||||
)
|
||||
raise SubscriptionRequiredError(
|
||||
message=message, subscribe_url=subscribe_url
|
||||
) from e
|
||||
|
||||
# Raise generic CloudAPIError with status code and detail
|
||||
raise CloudAPIError(
|
||||
f"API request failed: {e}",
|
||||
status_code=response.status_code,
|
||||
detail=error_detail if isinstance(error_detail, dict) else {},
|
||||
) from e
|
||||
except httpx.HTTPError as e:
|
||||
# Check if this is a response error with response details
|
||||
if hasattr(e, "response") and e.response is not None: # pyright: ignore [reportAttributeAccessIssue]
|
||||
response = e.response # type: ignore
|
||||
|
||||
# Try to parse error detail from response
|
||||
error_detail = None
|
||||
try:
|
||||
error_detail = response.json()
|
||||
except Exception:
|
||||
# If JSON parsing fails, we'll handle it as a generic error
|
||||
pass
|
||||
|
||||
# Check for subscription_required error (403)
|
||||
if response.status_code == 403 and isinstance(error_detail, dict):
|
||||
# Handle both FastAPI HTTPException format (nested under "detail")
|
||||
# and direct format
|
||||
detail_obj = error_detail.get("detail", error_detail)
|
||||
if (
|
||||
isinstance(detail_obj, dict)
|
||||
and detail_obj.get("error") == "subscription_required"
|
||||
):
|
||||
message = detail_obj.get("message", "Active subscription required")
|
||||
subscribe_url = detail_obj.get(
|
||||
"subscribe_url", "https://basicmemory.com/subscribe"
|
||||
)
|
||||
raise SubscriptionRequiredError(
|
||||
message=message, subscribe_url=subscribe_url
|
||||
) from e
|
||||
|
||||
# Raise generic CloudAPIError with status code and detail
|
||||
raise CloudAPIError(
|
||||
f"API request failed: {e}",
|
||||
status_code=response.status_code,
|
||||
detail=error_detail if isinstance(error_detail, dict) else {},
|
||||
) from e
|
||||
|
||||
raise CloudAPIError(f"API request failed: {e}") from e
|
||||
|
||||
@@ -345,7 +345,7 @@ def recent_activity(
|
||||
with force_routing(local=local, cloud=cloud):
|
||||
result = run_with_cleanup(
|
||||
mcp_recent_activity(
|
||||
type=type, # pyright: ignore[reportArgumentType]
|
||||
type=type or "",
|
||||
depth=depth if depth is not None else 1,
|
||||
timeframe=timeframe if timeframe is not None else "7d",
|
||||
page=page,
|
||||
|
||||
@@ -8,7 +8,7 @@ from dataclasses import dataclass
|
||||
from datetime import datetime
|
||||
from enum import Enum
|
||||
from pathlib import Path
|
||||
from typing import Any, Dict, Literal, Optional, List, Tuple
|
||||
from typing import TYPE_CHECKING, Any, Dict, Literal, Optional, List, Tuple
|
||||
|
||||
from loguru import logger
|
||||
from pydantic import AliasChoices, BaseModel, Field, model_validator
|
||||
@@ -122,6 +122,11 @@ class ProjectEntry(BaseModel):
|
||||
class BasicMemoryConfig(BaseSettings):
|
||||
"""Pydantic model for Basic Memory global configuration."""
|
||||
|
||||
if TYPE_CHECKING:
|
||||
# Pydantic accepts raw constructor data and validates/coerces it at runtime.
|
||||
# Model attributes remain strongly typed after initialization.
|
||||
def __init__(self, **data: Any) -> None: ...
|
||||
|
||||
env: Environment = Field(default="dev", description="Environment name")
|
||||
|
||||
projects: Dict[str, ProjectEntry] = Field(
|
||||
|
||||
@@ -39,23 +39,24 @@ def format_timestamp(timestamp: Any) -> str: # pragma: no cover
|
||||
Returns:
|
||||
A formatted string representation of the timestamp.
|
||||
"""
|
||||
parsed_timestamp = timestamp
|
||||
if isinstance(timestamp, str):
|
||||
try:
|
||||
# Try ISO format
|
||||
timestamp = datetime.fromisoformat(timestamp.replace("Z", "+00:00"))
|
||||
parsed_timestamp = datetime.fromisoformat(timestamp.replace("Z", "+00:00"))
|
||||
except ValueError:
|
||||
try:
|
||||
# Try unix timestamp as string
|
||||
timestamp = datetime.fromtimestamp(float(timestamp)).astimezone()
|
||||
parsed_timestamp = datetime.fromtimestamp(float(timestamp)).astimezone()
|
||||
except ValueError:
|
||||
# Return as is if we can't parse it
|
||||
return timestamp
|
||||
elif isinstance(timestamp, (int, float)):
|
||||
# Unix timestamp
|
||||
timestamp = datetime.fromtimestamp(timestamp).astimezone()
|
||||
parsed_timestamp = datetime.fromtimestamp(timestamp).astimezone()
|
||||
|
||||
if isinstance(timestamp, datetime):
|
||||
return timestamp.strftime("%Y-%m-%d %H:%M:%S")
|
||||
if isinstance(parsed_timestamp, datetime):
|
||||
return parsed_timestamp.strftime("%Y-%m-%d %H:%M:%S")
|
||||
|
||||
# Return as is if we can't format it
|
||||
return str(timestamp) # pragma: no cover
|
||||
return str(parsed_timestamp) # pragma: no cover
|
||||
|
||||
@@ -1,9 +1,9 @@
|
||||
"""Schema models for entity markdown files."""
|
||||
|
||||
from datetime import datetime
|
||||
from typing import List, Optional
|
||||
from typing import TYPE_CHECKING, Any, List, Optional
|
||||
|
||||
from pydantic import BaseModel
|
||||
from pydantic import BaseModel, Field, model_validator
|
||||
|
||||
|
||||
class Observation(BaseModel):
|
||||
@@ -38,23 +38,47 @@ class Relation(BaseModel):
|
||||
class EntityFrontmatter(BaseModel):
|
||||
"""Required frontmatter fields for an entity."""
|
||||
|
||||
metadata: dict = {}
|
||||
if TYPE_CHECKING:
|
||||
# Frontmatter may be built from raw YAML keys. The validator below
|
||||
# gathers those keys into the metadata mapping used at runtime.
|
||||
def __init__(self, **data: Any) -> None: ...
|
||||
|
||||
metadata: dict[str, Any] = Field(default_factory=dict)
|
||||
|
||||
@model_validator(mode="before")
|
||||
@classmethod
|
||||
def collect_metadata(cls, data: Any) -> Any:
|
||||
if not isinstance(data, dict):
|
||||
return data
|
||||
|
||||
if "metadata" not in data:
|
||||
return {"metadata": data}
|
||||
|
||||
metadata = data.get("metadata") or {}
|
||||
extras = {key: value for key, value in data.items() if key != "metadata"}
|
||||
if extras:
|
||||
return {"metadata": {**extras, **metadata}}
|
||||
return data
|
||||
|
||||
@property
|
||||
def tags(self) -> List[str]:
|
||||
return self.metadata.get("tags") if self.metadata else None # pyright: ignore
|
||||
tags = self.metadata.get("tags")
|
||||
return [str(tag) for tag in tags] if isinstance(tags, list) else []
|
||||
|
||||
@property
|
||||
def title(self) -> str:
|
||||
return self.metadata.get("title") if self.metadata else None # pyright: ignore
|
||||
title = self.metadata.get("title")
|
||||
return title if isinstance(title, str) else ""
|
||||
|
||||
@property
|
||||
def type(self) -> str:
|
||||
return self.metadata.get("type", "note") if self.metadata else "note" # pyright: ignore
|
||||
note_type = self.metadata.get("type", "note")
|
||||
return note_type if isinstance(note_type, str) else "note"
|
||||
|
||||
@property
|
||||
def permalink(self) -> str:
|
||||
return self.metadata.get("permalink") if self.metadata else None # pyright: ignore
|
||||
def permalink(self) -> Optional[str]:
|
||||
permalink = self.metadata.get("permalink")
|
||||
return permalink if isinstance(permalink, str) else None
|
||||
|
||||
|
||||
class EntityMarkdown(BaseModel):
|
||||
|
||||
@@ -95,8 +95,8 @@ def format_prompt_context(context: PromptContext) -> str:
|
||||
sections = []
|
||||
|
||||
# Process each context
|
||||
for context in context.results: # pyright: ignore
|
||||
for primary in context.primary_results: # pyright: ignore
|
||||
for context_item in context.results:
|
||||
for primary in context_item.primary_results:
|
||||
if primary.permalink not in added_permalinks:
|
||||
primary_permalink = primary.permalink
|
||||
|
||||
@@ -121,8 +121,8 @@ def format_prompt_context(context: PromptContext) -> str:
|
||||
section += f"- **Created**: {primary.created_at.strftime('%Y-%m-%d %H:%M')}\n"
|
||||
|
||||
# Add content snippet
|
||||
if hasattr(primary, "content") and primary.content: # pyright: ignore
|
||||
content = primary.content or "" # pyright: ignore # pragma: no cover
|
||||
if hasattr(primary, "content") and primary.content:
|
||||
content = primary.content or "" # pragma: no cover
|
||||
if content: # pragma: no cover
|
||||
section += f"\n**Excerpt**:\n{content}\n" # pragma: no cover
|
||||
|
||||
@@ -132,14 +132,14 @@ def format_prompt_context(context: PromptContext) -> str:
|
||||
""")
|
||||
sections.append(section)
|
||||
|
||||
if context.related_results: # pyright: ignore
|
||||
section += dedent( # pyright: ignore
|
||||
if context_item.related_results:
|
||||
section += dedent(
|
||||
"""
|
||||
## Related Context
|
||||
"""
|
||||
)
|
||||
|
||||
for related in context.related_results: # pyright: ignore
|
||||
for related in context_item.related_results:
|
||||
section_content = dedent(f"""
|
||||
- type: **{related.type}**
|
||||
- title: {related.title}
|
||||
|
||||
@@ -1,7 +1,7 @@
|
||||
"""Read note tool for Basic Memory MCP server."""
|
||||
|
||||
from textwrap import dedent
|
||||
from typing import Optional, Literal
|
||||
from typing import Optional, Literal, cast
|
||||
|
||||
import yaml
|
||||
|
||||
@@ -235,13 +235,22 @@ async def read_note(
|
||||
"frontmatter": None,
|
||||
}
|
||||
|
||||
def _search_results(payload: object) -> list[dict]:
|
||||
def _search_results(payload: object) -> list[dict[str, object]]:
|
||||
if not isinstance(payload, dict):
|
||||
return []
|
||||
results = payload.get("results")
|
||||
return results if isinstance(results, list) else []
|
||||
payload_dict = cast(dict[str, object], payload)
|
||||
results = payload_dict.get("results")
|
||||
if not isinstance(results, list):
|
||||
return []
|
||||
return [
|
||||
cast(dict[str, object], result)
|
||||
for result in results
|
||||
if isinstance(result, dict)
|
||||
]
|
||||
|
||||
async def _search_candidates(identifier_text: str, *, title_only: bool) -> dict:
|
||||
async def _search_candidates(
|
||||
identifier_text: str, *, title_only: bool
|
||||
) -> dict[str, object]:
|
||||
# Trigger: direct entity resolution failed for the caller's identifier.
|
||||
# Why: search_notes applies the same memory:// normalization and tool-level
|
||||
# query handling as the rest of MCP routing, which raw client calls skip.
|
||||
@@ -257,16 +266,16 @@ async def read_note(
|
||||
output_format="json",
|
||||
context=context,
|
||||
)
|
||||
return response if isinstance(response, dict) else {}
|
||||
return cast(dict[str, object], response) if isinstance(response, dict) else {}
|
||||
|
||||
def _result_title(item: dict) -> str:
|
||||
def _result_title(item: dict[str, object]) -> str:
|
||||
return str(item.get("title") or "")
|
||||
|
||||
def _result_permalink(item: dict) -> Optional[str]:
|
||||
def _result_permalink(item: dict[str, object]) -> Optional[str]:
|
||||
value = item.get("permalink")
|
||||
return str(value) if value else None
|
||||
|
||||
def _result_file_path(item: dict) -> Optional[str]:
|
||||
def _result_file_path(item: dict[str, object]) -> Optional[str]:
|
||||
value = item.get("file_path")
|
||||
return str(value) if value else None
|
||||
|
||||
|
||||
@@ -1,5 +1,7 @@
|
||||
"""Base model class for SQLAlchemy models."""
|
||||
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
from sqlalchemy.ext.asyncio import AsyncAttrs
|
||||
from sqlalchemy.orm import DeclarativeBase
|
||||
|
||||
@@ -7,4 +9,5 @@ from sqlalchemy.orm import DeclarativeBase
|
||||
class Base(AsyncAttrs, DeclarativeBase):
|
||||
"""Base class for all models"""
|
||||
|
||||
pass
|
||||
if TYPE_CHECKING:
|
||||
id: int
|
||||
|
||||
@@ -11,7 +11,7 @@ from basic_memory.repository.embedding_provider import EmbeddingProvider
|
||||
from basic_memory.repository.semantic_errors import SemanticDependenciesMissingError
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from fastembed import TextEmbedding # type: ignore[import-not-found] # pragma: no cover
|
||||
from fastembed import TextEmbedding # pragma: no cover
|
||||
|
||||
|
||||
class FastEmbedEmbeddingProvider(EmbeddingProvider):
|
||||
@@ -62,7 +62,7 @@ class FastEmbedEmbeddingProvider(EmbeddingProvider):
|
||||
|
||||
def _create_model() -> "TextEmbedding":
|
||||
try:
|
||||
from fastembed import TextEmbedding # type: ignore[import-not-found]
|
||||
from fastembed import TextEmbedding
|
||||
except (
|
||||
ImportError
|
||||
) as exc: # pragma: no cover - exercised via tests with monkeypatch
|
||||
|
||||
@@ -50,7 +50,7 @@ class OpenAIEmbeddingProvider(EmbeddingProvider):
|
||||
return self._client
|
||||
|
||||
try:
|
||||
from openai import AsyncOpenAI # type: ignore[import-not-found]
|
||||
from openai import AsyncOpenAI
|
||||
except ImportError as exc: # pragma: no cover - covered via monkeypatch tests
|
||||
raise SemanticDependenciesMissingError(
|
||||
"OpenAI dependency is missing. "
|
||||
|
||||
@@ -268,7 +268,7 @@ class Repository[T: Base]:
|
||||
|
||||
return await self.select_by_ids(session, [model.id for model in model_list]) # pyright: ignore [reportAttributeAccessIssue]
|
||||
|
||||
async def update(self, entity_id: int, entity_data: dict | T) -> Optional[T]:
|
||||
async def update(self, entity_id: int, entity_data: dict[str, Any] | T) -> Optional[T]:
|
||||
"""Update an entity with the given data."""
|
||||
logger.debug(f"Updating {self.Model.__name__} {entity_id} with data: {entity_data}")
|
||||
async with db.scoped_session(self.session_maker) as session:
|
||||
@@ -279,12 +279,13 @@ class Repository[T: Base]:
|
||||
entity = result.scalars().one()
|
||||
|
||||
if isinstance(entity_data, dict):
|
||||
for key, value in entity_data.items():
|
||||
update_data = cast(dict[str, Any], entity_data)
|
||||
for key, value in update_data.items():
|
||||
if key in self.valid_columns:
|
||||
setattr(entity, key, value)
|
||||
|
||||
elif isinstance(entity_data, self.Model):
|
||||
for column in self.Model.__table__.columns.keys():
|
||||
for column in self.valid_columns:
|
||||
setattr(entity, column, getattr(entity_data, column))
|
||||
|
||||
await session.flush() # Make sure changes are flushed
|
||||
|
||||
@@ -1068,39 +1068,48 @@ class SearchRepositoryBase(ABC):
|
||||
write_seconds_total=result.write_seconds_total,
|
||||
)
|
||||
batch_total_seconds = time.perf_counter() - batch_start
|
||||
metric_attrs = {
|
||||
"backend": backend_name,
|
||||
"skip_only_batch": result.embedding_jobs_total == 0,
|
||||
}
|
||||
telemetry.record_histogram(
|
||||
"vector_sync_batch_total_seconds",
|
||||
batch_total_seconds,
|
||||
unit="s",
|
||||
**metric_attrs,
|
||||
backend=backend_name,
|
||||
skip_only_batch=result.embedding_jobs_total == 0,
|
||||
)
|
||||
telemetry.add_counter(
|
||||
"vector_sync_entities_total", result.entities_total, **metric_attrs
|
||||
"vector_sync_entities_total",
|
||||
result.entities_total,
|
||||
backend=backend_name,
|
||||
skip_only_batch=result.embedding_jobs_total == 0,
|
||||
)
|
||||
telemetry.add_counter(
|
||||
"vector_sync_entities_skipped",
|
||||
result.entities_skipped,
|
||||
**metric_attrs,
|
||||
backend=backend_name,
|
||||
skip_only_batch=result.embedding_jobs_total == 0,
|
||||
)
|
||||
telemetry.add_counter(
|
||||
"vector_sync_entities_deferred",
|
||||
result.entities_deferred,
|
||||
**metric_attrs,
|
||||
backend=backend_name,
|
||||
skip_only_batch=result.embedding_jobs_total == 0,
|
||||
)
|
||||
telemetry.add_counter(
|
||||
"vector_sync_embedding_jobs_total",
|
||||
result.embedding_jobs_total,
|
||||
**metric_attrs,
|
||||
backend=backend_name,
|
||||
skip_only_batch=result.embedding_jobs_total == 0,
|
||||
)
|
||||
telemetry.add_counter(
|
||||
"vector_sync_chunks_total",
|
||||
result.chunks_total,
|
||||
backend=backend_name,
|
||||
skip_only_batch=result.embedding_jobs_total == 0,
|
||||
)
|
||||
telemetry.add_counter("vector_sync_chunks_total", result.chunks_total, **metric_attrs)
|
||||
telemetry.add_counter(
|
||||
"vector_sync_chunks_skipped",
|
||||
result.chunks_skipped,
|
||||
**metric_attrs,
|
||||
backend=backend_name,
|
||||
skip_only_batch=result.embedding_jobs_total == 0,
|
||||
)
|
||||
if batch_span is not None:
|
||||
batch_span.set_attributes(
|
||||
@@ -1675,33 +1684,33 @@ class SearchRepositoryBase(ABC):
|
||||
) -> None:
|
||||
"""Log completion and slow-entity warnings with a consistent format."""
|
||||
backend_name = type(self).__name__.removesuffix("SearchRepository").lower()
|
||||
metric_attrs = {
|
||||
"backend": backend_name,
|
||||
"skip_only_entity": entity_skipped and embedding_jobs_count == 0,
|
||||
}
|
||||
telemetry.record_histogram(
|
||||
"vector_sync_prepare_seconds",
|
||||
prepare_seconds,
|
||||
unit="s",
|
||||
**metric_attrs,
|
||||
backend=backend_name,
|
||||
skip_only_entity=entity_skipped and embedding_jobs_count == 0,
|
||||
)
|
||||
telemetry.record_histogram(
|
||||
"vector_sync_queue_wait_seconds",
|
||||
queue_wait_seconds,
|
||||
unit="s",
|
||||
**metric_attrs,
|
||||
backend=backend_name,
|
||||
skip_only_entity=entity_skipped and embedding_jobs_count == 0,
|
||||
)
|
||||
telemetry.record_histogram(
|
||||
"vector_sync_embed_seconds",
|
||||
embed_seconds,
|
||||
unit="s",
|
||||
**metric_attrs,
|
||||
backend=backend_name,
|
||||
skip_only_entity=entity_skipped and embedding_jobs_count == 0,
|
||||
)
|
||||
telemetry.record_histogram(
|
||||
"vector_sync_write_seconds",
|
||||
write_seconds,
|
||||
unit="s",
|
||||
**metric_attrs,
|
||||
backend=backend_name,
|
||||
skip_only_entity=entity_skipped and embedding_jobs_count == 0,
|
||||
)
|
||||
if total_seconds > 10:
|
||||
logger.warning(
|
||||
|
||||
@@ -350,7 +350,7 @@ class SQLiteSearchRepository(SearchRepositoryBase):
|
||||
pass
|
||||
|
||||
try:
|
||||
import sqlite_vec # type: ignore[import-not-found]
|
||||
import sqlite_vec
|
||||
except ImportError as exc:
|
||||
raise SemanticDependenciesMissingError(
|
||||
"sqlite-vec package is missing. "
|
||||
|
||||
@@ -103,7 +103,7 @@ MemoryUrl = Annotated[
|
||||
memory_url = TypeAdapter(MemoryUrl)
|
||||
|
||||
|
||||
def memory_url_path(url: memory_url) -> str: # pyright: ignore
|
||||
def memory_url_path(url: str) -> str:
|
||||
"""
|
||||
Returns the uri for a url value by removing the prefix "memory://" from a given MemoryUrl.
|
||||
|
||||
|
||||
@@ -194,7 +194,7 @@ class EntityResponse(SQLAlchemyModel):
|
||||
note_type: NoteType
|
||||
|
||||
# COMPAT(v0.18): old clients expect entity_type; remove when no longer needed
|
||||
@computed_field # type: ignore[prop-decorator]
|
||||
@computed_field
|
||||
@property
|
||||
def entity_type(self) -> str:
|
||||
return self.note_type
|
||||
|
||||
@@ -288,7 +288,9 @@ class EntityService(BaseService[EntityModel]):
|
||||
|
||||
if "permalink" in content_frontmatter:
|
||||
content_markdown = self._build_frontmatter_markdown(
|
||||
schema.title, schema.note_type, content_frontmatter["permalink"]
|
||||
schema.title,
|
||||
schema.note_type,
|
||||
_coerce_to_string(content_frontmatter["permalink"]),
|
||||
)
|
||||
|
||||
# Get unique permalink (prioritizing content frontmatter) unless disabled
|
||||
@@ -393,7 +395,9 @@ class EntityService(BaseService[EntityModel]):
|
||||
|
||||
if "permalink" in content_frontmatter:
|
||||
content_markdown = self._build_frontmatter_markdown(
|
||||
schema.title, schema.note_type, content_frontmatter["permalink"]
|
||||
schema.title,
|
||||
schema.note_type,
|
||||
_coerce_to_string(content_frontmatter["permalink"]),
|
||||
)
|
||||
|
||||
# Check if we need to update the permalink based on content frontmatter (unless disabled)
|
||||
@@ -522,7 +526,9 @@ class EntityService(BaseService[EntityModel]):
|
||||
|
||||
if "permalink" in content_frontmatter:
|
||||
content_markdown = self._build_frontmatter_markdown(
|
||||
schema.title, schema.note_type, content_frontmatter["permalink"]
|
||||
schema.title,
|
||||
schema.note_type,
|
||||
_coerce_to_string(content_frontmatter["permalink"]),
|
||||
)
|
||||
|
||||
# --- Permalink Resolution ---
|
||||
@@ -663,9 +669,9 @@ class EntityService(BaseService[EntityModel]):
|
||||
|
||||
if "permalink" in content_frontmatter:
|
||||
content_markdown = self._build_frontmatter_markdown(
|
||||
update_data.get("title", entity.title),
|
||||
update_data.get("note_type", entity.note_type),
|
||||
content_frontmatter["permalink"],
|
||||
_coerce_to_string(update_data.get("title", entity.title)),
|
||||
_coerce_to_string(update_data.get("note_type", entity.note_type)),
|
||||
_coerce_to_string(content_frontmatter["permalink"]),
|
||||
)
|
||||
|
||||
metadata = normalize_frontmatter_metadata(content_frontmatter or {})
|
||||
@@ -1002,7 +1008,7 @@ class EntityService(BaseService[EntityModel]):
|
||||
target_entity: Optional[Entity] = None
|
||||
if not isinstance(resolved, Exception):
|
||||
# Type narrowing: resolved is Optional[Entity] here, not Exception
|
||||
target_entity = resolved # type: ignore
|
||||
target_entity = resolved
|
||||
|
||||
# if the target is found, store the id
|
||||
target_id = target_entity.id if target_entity else None
|
||||
|
||||
@@ -149,7 +149,7 @@ class WatchService:
|
||||
|
||||
# create coroutines to handle changes
|
||||
change_handlers = [
|
||||
self.handle_changes(project, changes) # pyright: ignore
|
||||
self.handle_changes(project, set(changes))
|
||||
for project, changes in project_changes.items()
|
||||
]
|
||||
|
||||
@@ -502,19 +502,19 @@ class WatchService:
|
||||
|
||||
# Add a concise summary instead of a divider
|
||||
if processed:
|
||||
changes = [] # pyright: ignore
|
||||
change_summary: list[str] = []
|
||||
if add_count > 0:
|
||||
changes.append(f"[green]{add_count} added[/green]") # pyright: ignore
|
||||
change_summary.append(f"[green]{add_count} added[/green]")
|
||||
if modify_count > 0:
|
||||
changes.append(f"[yellow]{modify_count} modified[/yellow]") # pyright: ignore
|
||||
change_summary.append(f"[yellow]{modify_count} modified[/yellow]")
|
||||
if moved_count > 0:
|
||||
changes.append(f"[blue]{moved_count} moved[/blue]") # pyright: ignore
|
||||
change_summary.append(f"[blue]{moved_count} moved[/blue]")
|
||||
if delete_count > 0:
|
||||
changes.append(f"[red]{delete_count} deleted[/red]") # pyright: ignore
|
||||
change_summary.append(f"[red]{delete_count} deleted[/red]")
|
||||
|
||||
if changes:
|
||||
self.console.print(f"{', '.join(changes)}", style="dim") # pyright: ignore
|
||||
logger.info(f"changes: {len(changes)}")
|
||||
if change_summary:
|
||||
self.console.print(f"{', '.join(change_summary)}", style="dim")
|
||||
logger.info(f"changes: {len(change_summary)}")
|
||||
|
||||
duration_ms = int((time.time() - start_time) * 1000)
|
||||
self.state.last_scan = datetime.now()
|
||||
|
||||
Reference in New Issue
Block a user