diff --git a/lollms/functions/markdown2latex.py b/lollms/functions/markdown2latex.py index 8ed29ca..3a7ef31 100644 --- a/lollms/functions/markdown2latex.py +++ b/lollms/functions/markdown2latex.py @@ -28,7 +28,8 @@ def markdown_to_latex(file_path: str) -> str: # Define conversion rules from markdown to LaTeX conversion_rules = [ (r'\\', r'\\textbackslash{}'), # Escape backslashes - (r'([#]+) (.*)', lambda m: '\\' + 'sub'*(len(m.group(1))-1) + 'section{' + m.group(2) + '}'), # Convert headings + (r'^## (Chapter\s*\d+:\s*)?(.*)', lambda match: f"\\chapter{{{match.group(2)}}}"), # Convert chapters + (r'^### (.*)', r'\\section*{\1}'), # Convert subsections without numbering (r'\*\*(.*?)\*\*', r'\\textbf{\1}'), # Bold text (r'\*(.*?)\*', r'\\textit{\1}'), # Italic text (r'\!\[(.*?)\]\((.*?)\)', r'\\begin{figure}[h!]\n\\centering\n\\includegraphics[width=\\textwidth]{\2}\n\\caption{\1}\n\\end{figure}'), # Images @@ -43,14 +44,26 @@ def markdown_to_latex(file_path: str) -> str: # Apply conversion rules latex_text = markdown_text for pattern, replacement in conversion_rules: - latex_text = re.sub(pattern, replacement, latex_text, flags=re.MULTILINE) + if callable(replacement): + latex_text = re.sub(pattern, replacement, latex_text, flags=re.MULTILINE) + else: + latex_text = re.sub(pattern, replacement, latex_text, flags=re.MULTILINE) - # Create the LaTeX document structure + # Add lettrine to the first letter of the first paragraph of each chapter + latex_text = re.sub(r'(\\chapter{[^}]*}\n\n?)(\w)', r'\1\\lettrine[lines=2]{\2}{}', latex_text) + + # Create the LaTeX document structure with fancyhdr for headers latex_document = f""" \\documentclass{{book}} \\usepackage{{hyperref}} \\usepackage{{graphicx}} \\usepackage{{verbatim}} + \\usepackage{{fancyhdr}} + \\usepackage{{lettrine}} + \\pagestyle{{fancy}} + \\fancyhf{{}} + \\fancyhead[LE,RO]{{\\thepage}} + \\fancyhead[RE,LO]{{\\leftmark}} \\begin{{document}} diff --git a/lollms/functions/story_writing.py b/lollms/functions/story_writing.py index cd47b39..645b526 100644 --- a/lollms/functions/story_writing.py +++ b/lollms/functions/story_writing.py @@ -20,11 +20,53 @@ import json # Import markdown2latex from lollms.functions.markdown2latex import markdown_to_latex from lollms.functions.file_manipulation import change_file_extension +from lollms.functions.generate_image import build_image +from lollms.client_session import Client import copy +def needs_illustration(section_name, mydict): + for section in mydict['sections']: + if section['section_name'] == section_name: + return section['build_illustration'] + return False # Return None if the section is not found + + +def build_section_illustration(llm, prompt_ideas, current_section, content, client:Client): + discussion_prompt_separator = llm.config.discussion_prompt_separator + start_header_id_template = llm.config.start_header_id_template + end_header_id_template = llm.config.end_header_id_template + separator_template = llm.config.separator_template + system_message_template = llm.config.system_message_template + + prompt = "\n".join([ + f"{start_header_id_template}{system_message_template}{end_header_id_template}", + f"You are illustration builder for a book section. Your task is to build a prompt to generate an illustration from a book section of a story.", + "Your output must strictly contain only a json formatted output like this:", + "```json", + "{", + '"prompt":"image generation prompt",' + '"width":the width of the image,' + '"height":the height of the image' + "}" + "```", + f"{start_header_id_template}General idea{end_header_id_template}", + f"{prompt_ideas}\n\n" + f"{start_header_id_template}Section title{end_header_id_template}", + f"{current_section}", + f"{start_header_id_template}Section content{end_header_id_template}", + f"{content}", + f"{start_header_id_template}story_section_illustrator{end_header_id_template}" + ] + ) + image_generation_prompt = llm.fast_gen(prompt, callback=llm.sink) + code_blocks = llm.extract_code_blocks(image_generation_prompt) + if len(code_blocks)>0: + code = json.loads(code_blocks[0]["content"]) + return build_image(image_generation_prompt,width=code["width"],height=code["height"], client=client) + # Define the core functions -def start_writing_story(prompt_ideas: str, llm: Any, story_file_path: str, build_latex:bool=False) -> str: +def start_writing_story(prompt_ideas: str, llm: Any, story_file_path: str, build_latex:bool=False, client:Client = None) -> str: discussion_prompt_separator = llm.config.discussion_prompt_separator start_header_id_template = llm.config.start_header_id_template end_header_id_template = llm.config.end_header_id_template @@ -44,6 +86,7 @@ def start_writing_story(prompt_ideas: str, llm: Any, story_file_path: str, build "{\n" "\"section_name\": \"The name of the section\",\n" "\"section_description\": \"A short description of the section\"\n" + "\"build_illustration\": A boolean indicating if an illustration is needed for this section\n" "}, ...\n" "]\n" "}\n" @@ -79,7 +122,9 @@ def start_writing_story(prompt_ideas: str, llm: Any, story_file_path: str, build story_file_path=story_file_path, story_plan=story_plan, current_section=section_name, - prompt_ideas=prompt_ideas + prompt_ideas=prompt_ideas, + add_illustration=needs_illustration(section_name, story_plan), + client=client ) section_full["content"]=new_section final_story_content += f"\n## {section_name}\n\n{new_section}\n" @@ -116,10 +161,10 @@ def start_writing_story(prompt_ideas: str, llm: Any, story_file_path: str, build return trace_exception(e) # Define the metadata functions -def start_writing_story_function(llm, story_file_path, build_latex:bool=False): +def start_writing_story_function(llm, story_file_path, build_latex:bool=False, client:Client=None): return { "function_name": "start_writing_story", - "function": partial(start_writing_story, llm=llm, story_file_path=story_file_path, build_latex=build_latex), + "function": partial(start_writing_story, llm=llm, story_file_path=story_file_path, build_latex=build_latex, client=client), "function_description": "Starts writing a story based on the provided prompt ideas, generating a plan in JSON format, and writing the story section by section.", "function_parameters": [ {"name": "prompt_ideas", "type": "str"} @@ -128,7 +173,7 @@ def start_writing_story_function(llm, story_file_path, build_latex:bool=False): # Define the core function -def write_story_section(prompt_ideas: str, llm: Any, story_file_path: str, story_plan: dict, current_section: str) -> str: +def write_story_section(prompt_ideas: str, llm: Any, story_file_path: str, story_plan: dict, current_section: str, add_illustration:bool=False, client:Client=None) -> str: discussion_prompt_separator = llm.config.discussion_prompt_separator start_header_id_template = llm.config.start_header_id_template end_header_id_template = llm.config.end_header_id_template @@ -162,8 +207,11 @@ def write_story_section(prompt_ideas: str, llm: Any, story_file_path: str, story f"{start_header_id_template}story_section_writer{end_header_id_template}" ] ) - new_section += llm.fast_gen(prompt, callback=llm.sink) - + content = llm.fast_gen(prompt, callback=llm.sink) + new_section += content + if add_illustration and client: + illustration = build_section_illustration(llm, prompt_ideas, current_section, content, client) + new_section += "\n" + illustration # Write the new section to the story file story_path.write_text(new_section) @@ -194,6 +242,10 @@ def write_story_section(prompt_ideas: str, llm: Any, story_file_path: str, story new_section = f"## {current_section}\n\n" new_section += llm.fast_gen(prompt, callback=llm.sink).strip() + if add_illustration and client: + illustration = build_section_illustration(llm, prompt_ideas, current_section, content, client) + new_section += "\n" + illustration + # Append the new section to the story file story_path.write_text(story_content + "\n" + new_section)