utils.py

import shutil
from pathlib import Path
from string import Template


doc_template = Template("""

---
File: $file_path
---

$file_content

""".strip())


choice_prompt = Template("""

The user has asked the following question: $question

The goal is get the user the 1 most relevant documentation file to answer the question.

Here is the tree structure of the documentation. Your task is to return the numeric id \
associated with the most relevant .md and .mdx file.

<tree>
$tree_structure
</tree>

Sample response: "1.3.2"
Top 1 file id:

""".strip())


def copy_search_results(search_results, dest_folder):
    """Copy files from search results to destination folder."""
    for item in search_results[0]:
        file_path = item['entity']['file_path']
        dest_path = Path(dest_folder) / file_path
        
        dest_path.parent.mkdir(parents=True, exist_ok=True)
        shutil.copy2(file_path, dest_path)


def create_documentation_string(file_ids, temp_folder):
    """Create documentation string from file IDs using the template."""
    documentation_parts = []
    
    for file_id in file_ids:
        # Find the corresponding file in the temp folder
        docs_path = Path(temp_folder) / "docs"
        for file_path in docs_path.rglob("*.md*"):
            if file_id + "." in str(file_path):
                try:
                    with open(file_path, 'r', encoding='utf-8') as f:
                        content = f.read()
                    
                    formatted_doc = doc_template.substitute(
                        file_path=str(file_path.relative_to(docs_path)),
                        file_content=content
                    )
                    documentation_parts.append(formatted_doc)
                    break
                except Exception as e:
                    print(f"Error reading file {file_path}: {e}")
    
    return "\n\n".join(documentation_parts)