Shrook21/Code-Assistant

1import re
2from langchain_core.documents import Document
3
4def split_code_by_function(doc: Document):
5    text = doc.page_content
6    matches = list(re.finditer(r"^def\s+\w+\(.*?\):", text, re.MULTILINE))
7    if not matches:
8        return [doc]
9    starts = [m.start() for m in matches] + [len(text)]
10    return [
11        Document(page_content=text[starts[i]:starts[i+1]].strip(), metadata=doc.metadata)
12        for i in range(len(starts) - 1)
13    ]
14