src/streamer/parse/attach_markdown.py

@@ -1,22 +1,30 @@
 from __future__ import annotations
 
+from pydantic import BaseModel
+
 from streamer.parse.shard import Shard, StreamFile
 
 
-class ShardWithMarkdown(Shard):
-    children: list[ShardWithMarkdown]
+class ShardWithMarkdown(BaseModel):
+    markers: list[str] = []
+    tags: list[str] = []
+    start_line: int
+    end_line: int
+    children: list["ShardWithMarkdown"] = []
     markdown_content: str
 
 
-class StreamFileWithMarkdown(StreamFile):
+class StreamFileWithMarkdown(BaseModel):
+    file_name: str
     shard: ShardWithMarkdown | None = None
 
 
 def attach_markdown_shard(shard: Shard, markdown_text: str) -> ShardWithMarkdown:
     lines = markdown_text.splitlines()
     markdown_content = "\n".join(lines[shard.start_line - 1 : shard.end_line])
+
     return ShardWithMarkdown(
-        **shard.model_dump(exclude=["children"]),
+        **shard.model_dump(exclude={"children"}),
         children=[
             attach_markdown_shard(child, markdown_text) for child in shard.children
         ],