split up entity parsing

This commit is contained in:
phernandez
2024-12-21 16:36:14 -06:00
parent 4f5a1c24e2
commit 7dfea0c704
3 changed files with 86 additions and 183 deletions
+3 -80
View File
@@ -3,12 +3,8 @@
import logging
from pathlib import Path
from markdown_it import MarkdownIt
from basic_memory.markdown.exceptions import ParseError
from basic_memory.markdown.schemas import (
Observation,
Relation,
Entity,
EntityFrontmatter,
EntityContent,
@@ -31,9 +27,6 @@ def debug_sections(text):
class EntityParser:
"""Parser for entity markdown files."""
def __init__(self):
self.md = MarkdownIt()
def parse_file(self, path: Path, encoding: str = "utf-8") -> Entity:
"""Parse an entity markdown file."""
if not path.exists():
@@ -50,80 +43,10 @@ class EntityParser:
# Parse each section using schema methods
frontmatter = EntityFrontmatter.from_text(sections[1])
content = EntityContent.from_markdown(sections[2])
metadata = EntityMetadata.from_text(sections[4] if len(sections) >= 5 else "")
# Parse markdown content (middle section)
content_tokens = self.md.parse(sections[2].strip())
# State for content parsing
title = ""
description = ""
observations = []
relations = []
current_section = None
# Track list items
in_list_item = False
list_item_tokens = []
for token in content_tokens:
if token.type == "heading_open":
if token.tag == "h1":
current_section = "title"
elif token.tag == "h2":
current_section = "section_name"
elif token.type == "inline":
content = token.content.strip()
if current_section == "title":
title = content
current_section = "description"
elif current_section == "section_name":
current_section = content.lower()
elif current_section == "description":
if description:
description += " "
description += content
elif in_list_item:
list_item_tokens.append(token)
elif token.type == "list_item_open":
in_list_item = True
list_item_tokens = []
elif token.type == "list_item_close":
item_content = " ".join(t.content for t in list_item_tokens)
try:
if current_section == "observations":
if obs := Observation.from_line(item_content):
observations.append(obs)
elif current_section == "relations":
if rel := Relation.from_line(item_content):
relations.append(rel)
except ParseError:
# Skip malformed items
pass
in_list_item = False
# Create content object
content = EntityContent(
title=title,
description=description,
observations=observations,
relations=relations,
)
# Parse metadata from final section
metadata_obj = EntityMetadata(metadata={})
if len(sections) >= 5:
metadata_text = sections[4].strip()
logger.debug(f"Metadata text: {metadata_text}")
for line in metadata_text.split("\n"):
if ":" in line:
key, value = line.split(":", 1)
metadata_obj.metadata[key.strip()] = value.strip()
return Entity(frontmatter=frontmatter, content=content, metadata=metadata_obj)
return Entity(frontmatter=frontmatter, content=content, metadata=metadata)
except UnicodeError as e:
if encoding == "utf-8":
@@ -4,6 +4,7 @@ import logging
from datetime import datetime
from typing import Any, Dict, List, Optional
from markdown_it import MarkdownIt
from pydantic import BaseModel
from basic_memory.markdown.exceptions import ParseError
@@ -47,6 +48,20 @@ class EntityMetadata(BaseModel):
metadata: Dict[str, Any] = {}
@classmethod
def from_text(cls, text: str) -> "EntityMetadata":
"""Parse metadata from text."""
try:
metadata = {}
if text: # Only parse if there's content
for line in text.strip().split("\n"):
if ":" in line:
key, value = line.split(":", 1)
metadata[key.strip()] = value.strip()
return cls(metadata=metadata)
except Exception as e:
raise ParseError(f"Failed to parse metadata: {e}") from e
class EntityContent(BaseModel):
"""Content sections of an entity markdown file."""
@@ -57,6 +72,74 @@ class EntityContent(BaseModel):
relations: List[Relation] = []
context: Optional[str] = None
@classmethod
def from_markdown(cls, text: str) -> "EntityContent":
"""Parse content from markdown text."""
try:
md = MarkdownIt()
tokens = md.parse(text.strip())
# State for parsing
title = ""
description = ""
observations: List[Observation] = []
relations: List[Relation] = []
current_section = None
# Track list items
in_list_item = False
list_item_tokens = []
for token in tokens:
if token.type == "heading_open":
if token.tag == "h1":
current_section = "title"
elif token.tag == "h2":
current_section = "section_name"
elif token.type == "inline":
content = token.content.strip()
if current_section == "title":
title = content
current_section = "description"
elif current_section == "section_name":
current_section = content.lower()
elif current_section == "description":
if description:
description += " "
description += content
elif in_list_item:
list_item_tokens.append(token)
elif token.type == "list_item_open":
in_list_item = True
list_item_tokens = []
elif token.type == "list_item_close":
item_content = " ".join(t.content for t in list_item_tokens)
try:
if current_section == "observations":
if obs := Observation.from_line(item_content):
observations.append(obs)
elif current_section == "relations":
if rel := Relation.from_line(item_content):
relations.append(rel)
except ParseError:
# Skip malformed items
pass
in_list_item = False
return cls(
title=title,
description=description,
observations=observations,
relations=relations,
)
except Exception as e:
raise ParseError(f"Failed to parse markdown content: {e}") from e
class Entity(BaseModel):
"""Complete entity combining frontmatter, content, and metadata."""