chore(core): use ty for typechecking

Signed-off-by: phernandez <paul@basicmachines.co>
This commit is contained in:
phernandez
2026-04-13 01:42:22 -05:00
parent abd4a5a6da
commit 58dd6963bd
88 changed files with 961 additions and 551 deletions
@@ -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,
+72 -45
View File
@@ -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
+1 -1
View File
@@ -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,
+6 -1
View File
@@ -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(
+7 -6
View File
@@ -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
+32 -8
View File
@@ -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):
+7 -7
View File
@@ -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}
+18 -9
View File
@@ -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
+4 -1
View File
@@ -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. "
+4 -3
View File
@@ -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. "
+1 -1
View File
@@ -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.
+1 -1
View File
@@ -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
+13 -7
View File
@@ -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
+9 -9
View File
@@ -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()