utils/helpers.py

from typing import Dict, Any
from llama_index.core import Response


def process_response(response: Response) -> Dict[str, Any]:
    source_nodes = response.source_nodes
    sources = {}
    for i, node in enumerate(source_nodes, 1):
        source = format_source(node.metadata)
        if source not in sources.values():
            sources[i] = source

    markdown_text = response.response + "\n\n### Sources\n\n"
    raw_text = response.response + "\n\nSources:\n"

    for i, source in sources.items():
        markdown_text += f"{i}. {source}\n"
        raw_text += f"[{i}] {source}\n"

    return {"markdown": markdown_text, "raw": raw_text, "sources": sources}


def format_source(metadata: Dict[str, Any]) -> str:
    authors = metadata.get("authors", "Unknown Author")
    year = metadata.get("year", "n.d.")
    title = metadata.get("title", "Untitled")

    author_list = authors.split(",")
    if len(author_list) > 2:
        formatted_authors = f"{author_list[0].strip()} et al."
    elif len(author_list) == 2:
        formatted_authors = f"{author_list[0].strip()} and {author_list[1].strip()}"
    else:
        formatted_authors = author_list[0].strip()

    year = "n.d." if year is None or year == "None" else str(year)

    max_title_length = 250
    if len(title) > max_title_length:
        title = title[:max_title_length] + "..."

    return f"{formatted_authors} ({year}). {title}"