"""**Scholaris** is developed using 'literate programming'. The sections below describe its core functionality, including helper functions for JSON schema generation, local file processing, external data retrieval from sources like NCBI, Semantic Scholar and OpenAlex, and an Assistant class for simplified chat and tool use."""

# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/01_core.ipynb.

# %% auto 0
__all__ = ['T', 'generate_json_schema', 'json_schema_decorator', 'get_file_names', 'extract_text_from_pdf',
           'extract_title_and_first_author', 'get_titles_and_first_authors', 'summarize_local_document',
           'describe_python_code', 'convert_id', 'detect_id_type', 'id_converter_tool', 'query_openalex_api',
           'query_semantic_scholar_api', 'respond_to_generic_queries', 'Assistant', 'add_to_class']

# %% ../nbs/01_core.ipynb 34
import inspect
from typing import Callable, Dict, List, Tuple, Optional, Any, Union 
import json
import os
import random
import ollama

# %% ../nbs/01_core.ipynb 35
def generate_json_schema(func: Callable) -> Dict[str, Any]:
    """
    Generate a JSON schema for the given function based on its annotations and docstring.
    
    Args:
        func (Callable): The function to generate a schema for.
    
    Returns:
        Dict[str, Any]: A JSON schema for the function.
    """
    annotations = func.__annotations__
    doc = inspect.getdoc(func)
    
    schema = {
        "type": "function",
        "function": {
            "name": func.__name__,
            "description": "",
            "parameters": {
                "type": "object",
                "properties": {},
                "required": []
            }
        }
    }
    
    if doc:
        lines = doc.split('\n')
        description = []
        arg_descriptions = {}
        in_args_section = False
        
        for line in lines:
            line = line.strip()
            if line.lower().startswith('args:'):
                in_args_section = True
                continue
            elif line.lower().startswith('returns:') or line.lower().startswith('raises:'):
                break
            
            if not in_args_section:
                description.append(line)
            else:
                parts = line.split(':')
                if len(parts) >= 2:
                    # print(parts) # Uncomment for debugging
                    arg_name = parts[0].split(' ')[0].strip()
                    arg_desc = ':'.join(parts[1:]).strip()
                    # print(arg_name, arg_desc) # Uncomment for debugging
                    arg_descriptions[arg_name] = arg_desc
                    # print(arg_descriptions) # Uncomment for debugging
                    if 'optional' not in parts[0].lower():
                        schema['function']['parameters']['required'].append(arg_name)

        
        schema['function']['description'] = ' '.join(description).strip()
    
    for arg, arg_type in annotations.items():
        if arg != 'return':
            schema['function']['parameters']['properties'][arg] = {
                "type": _get_type(arg_type),
                "description": arg_descriptions.get(arg, "")
            }
    
    return schema

def _get_type(arg_type):
    if 'List' in str(arg_type):
        return "list"
    elif 'Optional' in str(arg_type):
        return "optional"
    elif arg_type == int:
        return "integer"
    elif arg_type == str:
        return "string"
    elif arg_type == float:
        return "number"
    elif arg_type == bool:
        return "boolean"
    else:
        return "object"

# %% ../nbs/01_core.ipynb 36
import functools
from typing import TypeVar

# %% ../nbs/01_core.ipynb 37
T = TypeVar('T', bound=Callable)

def json_schema_decorator(func: T) -> T:
    """
    Decorator to generate and attach a JSON schema to a function.
    
    Args:
        func (Callable): The function to decorate.
    
    Returns:
        Callable: The decorated function with an attached JSON schema.
    """
    @functools.wraps(func)
    def wrapper(*args, **kwargs):
        return func(*args, **kwargs)
    
    schema = generate_json_schema(func)
    wrapper.json_schema = schema  # Attach the schema dictionary directly
    return wrapper

# %% ../nbs/01_core.ipynb 46
@json_schema_decorator
def get_file_names(ext: str = "pdf, txt") -> str:
    """Retrieves a list of file names with specified extensions in a local data directory the assistant has access to on the user's computer.

    Args:
        ext: A comma-separated string of file extensions to filter the files by. Options are: pdf, txt, md, markdown, csv, and py. Defaults to "pdf, txt".

    Returns:
        str: A comma-separated string of file names with the specified extensions. If no files are found, a message is returned.

    Example:
        >>> get_file_names(ext="pdf, txt")
        
        "List of file names with the specified extensions in the local data directory: file1.pdf, file2.txt"
    """

    if 'DIR_PATH' not in globals():
        return "Error: The local data directory path is not defined."
    
    if not os.path.exists(DIR_PATH):
        return f"Error: The local data directory does not exist."

    valid_extensions = ["pdf", "txt", "md", "markdown", "csv", "py"]

    # Process the input extensions
    selected_extensions = []
    for e in ext.split(','):
        e = e.lower().strip().lstrip('.').strip('{').strip('}').strip('[').strip(']').strip('(').strip(')').strip('"').strip("'") # Clean up the extension string
        if e not in valid_extensions:
            return f"Error: Invalid file extension '{e}'. Please choose from: pdf, txt, md, markdown, csv, py." # Instead of raising an error, we return a message to the LLM to avoid stopping a conversation
        selected_extensions.append(e)

    # List all files with the specified extensions
    file_names = [file for file in os.listdir(DIR_PATH) if any(file.endswith(f".{e}") for e in selected_extensions)]
    # print(file_names) # Uncomment for debugging

    if len(file_names) == 0:
         return "Access of local data directory successful but no files found with the specified extensions."

    file_names_json = ', '.join(file_names)

    # Convert list to a comma-separated string. This is because the object is returned to the LLM and the API accepts str only
    return f"List of file names with the specified extensions in the local data directory: {file_names_json}"

# %% ../nbs/01_core.ipynb 47
assert type(get_file_names.json_schema) == dict
assert get_file_names.json_schema['function']['name'] == "get_file_names"

# %% ../nbs/01_core.ipynb 58
import PyPDF2


# %% ../nbs/01_core.ipynb 59
@json_schema_decorator
def extract_text_from_pdf(file_name: str, page_range: Optional[str] = None) -> str:
    """A function that extracts text from a PDF file.
    Use this tool to extract specific details from a PDF document, such as abstract, authors, conclusions, or user-specified content of other sections.
    If the user specifies a page rage, use the optional page_range parameter to extract text from specific pages.
    If the user uses words such as beginning, middle, or end, to descripe the section, infer the page range based on the total number of 15 pages in a document.
    Do not use this tool to summarize an entire PDF document. Only use this tool for documents with extensions .pdf, or .PDF.

    Args:
        file_name (str): The file name of the PDF document in the local data directory.
        page_range (Optional[str]): A string with page numbers and/or page ranges separated by commas (e.g., "1" or "1, 2", or "5-7"). Default is None to extract all pages.

    Returns:
        str: Extracted text from the PDF.

    Example:
    >>> text = extract_text_from_pdf("./test.pdf", page_range="1")
    """
    verbose = False  # Set to True for debugging

    # Check if the directory exists
    if not os.path.exists(DIR_PATH):
        return "Local data directory not found."
    
    # Construct the full file path; this will be handled by a try-except block because this function is alos used as a tool
    try:
        pdf_path = os.path.join(DIR_PATH, file_name)
    except Exception as e:
        return f"Error: {e}"
    
    # Validate input
    if not os.path.exists(pdf_path):
        raise FileNotFoundError(f"File not found: {pdf_path}")
    assert pdf_path.endswith('.pdf'), "The file must be a PDF."
    assert isinstance(page_range, (str, type(None))), "The page_range must be a string or None."

    if page_range:
        # Clean up page range input in case of LLM formatting errors; must be a string, e.g., "1" or "1, 2", or "5-7", with no quotes or brackets   
        page_range = page_range.strip().replace('"', '').replace("'", "").replace("[", "").replace("]", "")

        # Parse the page range string
        page_numbers = []
        for part in page_range.split(','):
            if '-' in part:
                a, b = part.split('-')
                page_numbers.extend(range(int(a), int(b) + 1))
            else:
                page_numbers.append(int(part))
        if verbose:
            print(f"Extracting text from pages: {page_numbers}")

    # Validate page numbers
    if page_range:
        start_page, end_page = min(page_numbers), max(page_numbers)
        if start_page < 1 or end_page < start_page:
            raise ValueError("Invalid page range. Please provide a valid range of pages to extract text from.")

    # Extract text from the PDF
    text = ""
    try:
        with open(pdf_path, 'rb') as file:
            pdf = PyPDF2.PdfReader(file)
            
            # Check for invalid page numbers
            if page_range:
                max_page = len(pdf.pages)
                invalid_pages = [p for p in page_numbers if p < 1 or p > max_page]
                if invalid_pages:
                    page_range = None # Reset page range to extract all pages
            
            if page_range:
                for page_num in page_numbers:
                    if verbose:
                        print(f"Extracting text from page {page_num}...")
                    page = pdf.pages[page_num - 1]  # Adjust for 0-based index
                    text += f"page {page_num} of {len(pdf.pages)}\n{page.extract_text()}"
            else:
                for page_num, page in enumerate(pdf.pages):
                    if verbose:
                        print(f"Extracting text from page {page_num + 1}...")
                    text += f"page {page_num} of {len(pdf.pages)}\n{page.extract_text()}"
            
            if verbose:
                word_count = len(text.split())
                total_pages = len(page_numbers) if page_range else len(pdf.pages)
                print(f"Text extraction completed.\nTotal pages extracted: {total_pages}\nWord count: {word_count}\nNo. of characters (with spaces): {len(text)}")
    except PyPDF2.errors.PdfReadError as e:
        return f"Error reading PDF: {e}"

    return text

# %% ../nbs/01_core.ipynb 62
from tqdm import tqdm


# %% ../nbs/01_core.ipynb 63
def extract_title_and_first_author(contents: List[Dict[str, str]], model: str='llama3.1', verbose: Optional[bool] = False, show_progress: Optional[bool] = False) -> List[Dict[str, str]]:
    """
    A function that extracts the titles and the first author's names from the text of one or more research articles.

    Args:
        contents (List[Dict[str, str]]): A list of dictionaries containing the file name and extracted text.
        model (str): The model to use for the extraction. Default is 'llama3.1'.
        verbose (Optional[bool]): Whether to print additional information. Default is False.
        show_progress (Optional[bool]): Whether to show a progress bar. Default is False.

    Returns:
        contents (List[Dict[str, str]]): The input list of dictionaries with the extracted title and first author added.

    Raises:
        JSONDecodeError: If the JSON response is invalid.

    Example:
    >>> contents = extract_title_and_first_author(contents)
    Extracting titles and first authors: 100%|██████████| 3/3 [00:22<00:00,  7.35s/it]
    """

    prompt = """
    The text below between the <text> XML like tags is extracted from the first page of a research article. 
    Your task is to identify the title of the research article and the first author's name.
    The title is typically located immediately before the authors' names and the abstract.

    <text>
    {text}
    </text>

    The output must be provided in JSON format shown in the following example.

    Example output:
    {{
        "title": "<title>",
        "first_author": "<first_author>"
    }}
    Write the JSON output and nothing more. Do not include degree titles or affiliations in the author's name.

    Here is the JSON output:
    """

    pdf_iterator = tqdm(contents, desc="Extracting titles and first authors") if show_progress else contents
    
    for pdf_item in pdf_iterator:
        text = pdf_item['extracted_text']
        if verbose:
            tqdm.write(pdf_item['file_path'])
            tqdm.write("---")
            tqdm.write(text[:500])
            tqdm.write("---")

        try:
            response = ollama.chat(model=model, messages=[
                {
                    'role': 'user',
                    'content': prompt.format(text=text),
                },
            ])
        except Exception as e:
            tqdm.write(f"Error: {e}")
            if "Connection refused" in str(e):
                tqdm.write("Make sure Ollama is running and the correct model is available.")
            continue

        if verbose:
            tqdm.write(response['message']['content'])

        try:
            extracted_info = json.loads(response['message']['content'])
        except json.JSONDecodeError:
            tqdm.write(f"Error: Invalid JSON response for {pdf_item['file_path']}")
            extracted_info = {"title": "", "first_author": ""}

        pdf_item.update(extracted_info)

    print("\n") if show_progress else None # Add a newline if showing progress bar
    return contents

# %% ../nbs/01_core.ipynb 71
@json_schema_decorator
def get_titles_and_first_authors() -> str:
    """
    A function that retrieves the titles of research articles from a directory of PDF files.

    Returns:
        str: A JSON-formatted string containing the titles, first authors and file names of the research articles.

    Raises:
        FileNotFoundError: If the specified directory does not exist.

    Example:
    >>> get_titles_and_first_authors()
    """

    # Check if the directory exists
    if not os.path.exists(DIR_PATH):
        return "Local data directory not found."

    # Initialize an empty list to store the file names and extracted text
    pdf_contents = []

    # Initialize an empty list of dictionaries to store file names, extracted titles and first authors
    titles_and_authors = []

    # Get the file paths of all PDF files in the local data directory
    file_paths = [os.path.join(DIR_PATH, file) for file in os.listdir(DIR_PATH) if file.endswith('.pdf')]

    # Extract text from the first page of each PDF file
    for file_path in file_paths:
        pdf_contents.append({
            "file_name": file_path.split("/")[-1],
            "extracted_text": extract_text_from_pdf(file_path, page_range="1")
            })

    # Extract titles and first authors from the extracted text of each PDF file
    response = extract_title_and_first_author(pdf_contents, show_progress=True)
    for article in response:
        titles_and_authors.append({
            "title": article.get('title'), 
            "first_author": article.get('first_author'),
            "file_name": article.get('file_name'),
            })
    
    if not titles_and_authors:
        return "No titles found."
        
    return json.dumps(titles_and_authors, indent=2) # Return as a JSON-formatted string

# %% ../nbs/01_core.ipynb 73
assert type(get_titles_and_first_authors.json_schema) == dict
assert get_titles_and_first_authors.json_schema['function']['name'] == "get_titles_and_first_authors"

# %% ../nbs/01_core.ipynb 76
@json_schema_decorator
def summarize_local_document(file_name: str, ext: str = "pdf") -> str:
    """Summarize the content of a single PDF, markdown, or text document from the local data directory.

    Args:
        file_name (str): The file name of the local document to summarize.
        ext (str): The extension of the local document. Options are: pdf, txt, md, and markdown. Defaults to "pdf".

    Returns:
        str: The summary of the content of the local document.

    Example:
        >>> summarize_local_document("research_paper", ext="pdf")
    """
    # Check if the directory exists
    if not os.path.exists(DIR_PATH):
        return "Local data directory not found."

    # Ensure the extension is valid: delete spaces, hyphens, and quotation marks, and convert to lowercase in case of LLM input errors
    ext = ext.lower().replace('"', '').replace("'", "").replace(" ", "").replace(".", "")
    if ext not in ["pdf", "txt", "md", "markdown"]:
        return f"Invalid file extension '{ext}'. Please choose from: pdf, txt, md, markdown."

    # Get the file paths of all files in the local data directory
    file_paths = [os.path.join(DIR_PATH, file) for file in os.listdir(DIR_PATH) if file.endswith(ext)]
    # print(file_paths) # Uncomment for debugging
    
    # Find the file path that matches the specified file name
    file_path = [path for path in file_paths if file_name in path]
    # print(file_path) # Uncomment for debugging

    if not file_path:
        return f"No file found with the name '{file_name}' and extension '{ext}'."
    elif len(file_path) > 1:
        return f"Multiple files found with the name '{file_name}' and extension '{ext}'. Please specify a unique file name."
    else:
        file_path = file_path[0] # Convert the file path list to a string

    if ext == "pdf":
        # Extract text from a PDF file
        try:
            full_text = extract_text_from_pdf(file_path, page_range=None)
        except Exception as e:
            return f"Error while extracting text from PDF file {file_name}: {e}"
    
    if ext in ["txt", "md", "markdown"]:
        # Read the full text content from a text file
        try:
            with open(file_path, 'r') as file:
                full_text = file.read()
        except Exception as e:
            return f"Error while reading the file {file_name}: {e}"
    # print(full_text[:500]) # Uncomment for debugging

    # Remove references section from the full text content, if present
    patterns = ["References", "REFERENCES", "references", "Bibliography", "BIBLIOGRAPHY", "bibliography"]
    for pattern in patterns:
        if pattern in full_text:
            full_text = full_text.split(pattern)[0]
            break

    # Summarize the full text content of the document
    prompt = f"""
    The text below is the full text content from single document with the file name '{file_name}'.

    <text>
    {full_text}
    </text>

    If the document has an abstract, use the abstract for the summary. The abstract is typically located at the beginning of the document (page 1) and provides a concise summary of the research.
    If there is no abstract, generate a concise summary (approx. 200 words) that captures the main points of the document, including key findings and conclusions.
    Remember that your task is to summarize the content of the main text accurately and concisely, ignoring acknowledgements and references that are typically listed at the end of the document, after the conclusion section.
    Start the summary with the title of the document, typically found on the first page before the author names.
    """

    sys_message = """You are a scientific summarization assistant for health and life sciences research. 
    Your task is to condense the contents of a complex research document with mutiple pages into a concise, accurate summary.
    If the document describes a research study, highlight the main findings, methodologies, and conclusions of the study.
    Start the summary with the title of the document found on the first page, followed by a brief summary of the content.
    Ignore references typically listed at the end of a document after the conclusion section, as well as acknowledgements, and other non-content sections.
    """

    # Set the model from the global variables
    if 'MODEL' in globals():
        model = MODEL
    else:
        model = "llama3.1" # Default model

    try:
        response = ollama.chat(
            model=model,
            messages=[
                {'role': "system", 'content': sys_message},
                {'role': 'user', 'content': prompt},
            ]
        )
        # TODO: Try out different optional parameters for the ollama.chat function, such as temperature, max_tokens, etc. to improve the quality of the summary
        # For details, see the Ollama API documentation:
        # https://github.com/ollama/ollama/blob/main/docs/api.md
        # https://github.com/ollama/ollama/blob/main/docs/modelfile.md#valid-parameters-and-values

        summary = response['message']['content']
        return summary

    except Exception as e:
        return f"Error while summarizing the content of the document '{file_name}': {e}" 

# %% ../nbs/01_core.ipynb 78
assert type(summarize_local_document.json_schema) == dict
assert summarize_local_document.json_schema['function']['name'] == "summarize_local_document"
assert 'file_name' in summarize_local_document.json_schema['function']['parameters']['properties'].keys()
assert 'ext' in summarize_local_document.json_schema['function']['parameters']['properties'].keys()

# %% ../nbs/01_core.ipynb 82
@json_schema_decorator
def describe_python_code(file_name: str) -> str:
    """Describe the purpose of the Python code in a local Python file.
    This may involve summarizing the entire code, extracting key functions, or providing an overview of the code structure.

    Args:
        file_name (str): The file name of the local Python code file document to describe.

    Returns:
        str: A description of the purpose of the Python code in the local file.

    Example:
        >>> describe_python_code("main.py", ext="py")
    """
    # Check if the directory exists
    if not os.path.exists(DIR_PATH):
        return "Local data directory not found."

    # Get the file paths of all files in the local data directory
    file_paths = [os.path.join(DIR_PATH, file) for file in os.listdir(DIR_PATH) if file.endswith(".py")]
    # print(file_paths) # Uncomment for debugging
    
    # Find the file path that matches the specified file name
    file_path = [path for path in file_paths if file_name in path]
    # print(file_path) # Uncomment for debugging

    if not file_path:
        return f"No file found with the name '{file_name}' and extension '.py'."
    elif len(file_path) > 1:
        return f"Multiple files found with the name '{file_name}' and extension '.py'. Please specify a unique file name."
    else:
        file_path = file_path[0] # Convert the file path list to a string

    try:
        with open(file_path, 'r') as file:
            full_text = file.read()
    except Exception as e:
        return f"Error while reading the file {file_name}: {e}"
    # print(full_text[:500]) # Uncomment for debugging

    # Summarize the full text content of the document
    prompt = f"""
    The text below is the full Python code content from the file '{file_name}'.

    <code>
    {full_text}
    </code>

    Your task is to describe the purpose of the Python code in the file. This may involve summarizing the entire code, extracting key functions, or providing an overview of the code structure.
    """

    sys_message = """You are a programming assistant for Python code. Your task is to describe the purpose of Python code in a local file for the user who may not be familiar with the code,
    and may not know how to interpret the code. If not ask for specific content, provide a high-level overview of the code's functionality, key functions, and structure.
    """

    # Set the model from the global variables
    if 'MODEL' in globals():
        model = MODEL
    else:
        model = "llama3.1" # Default model

    try:
        response = ollama.chat(
            model=model,
            messages=[
                {'role': "system", 'content': sys_message},
                {'role': 'user', 'content': prompt},
            ]
        )
        # TODO: Try out different optional parameters for the ollama.chat function, such as temperature, max_tokens, etc. to improve the quality of the summary
        # For details, see the Ollama API documentation:
        # https://github.com/ollama/ollama/blob/main/docs/api.md
        # https://github.com/ollama/ollama/blob/main/docs/modelfile.md#valid-parameters-and-values

        summary = response['message']['content']
        return summary

    except Exception as e:
        return f"Error while describing the Python code in the file '{file_name}': {e}"

# %% ../nbs/01_core.ipynb 84
assert type(describe_python_code.json_schema) == dict
assert describe_python_code.json_schema['function']['name'] == "describe_python_code"
assert 'file_name' in describe_python_code.json_schema['function']['parameters']['properties'].keys()

# %% ../nbs/01_core.ipynb 89
import re
import requests

# %% ../nbs/01_core.ipynb 90
def convert_id(ids: List[str]) -> str:
    """
    For any article(s) in PubMed Central, find all the corresponding PubMed IDs (PMIDs), digital object identifiers (DOIs), and manuscript IDs (MIDs).
    
    Args:
    ids (List[str]): A list of IDs to convert (max 200 per request).
    
    Returns:
    Str: A JSON-formatted string containing the conversion results.
    """
    
    if 'EMAIL' in globals():
        email = EMAIL
    else:
        try:
            email = os.environ['EMAIL']
        except KeyError:
            return {"error": "Please provide an email address"}

    # API endpoint
    base_url = "https://www.ncbi.nlm.nih.gov/pmc/utils/idconv/v1.0/"
    
    # Prepare the IDs string
    ids_string = ",".join(ids)
    
    # Prepare the parameters
    params = {
        "tool": "scholaris",
        "email": email,
        "ids": ids_string,
        "format": "json"
    }
    
    try:
        # Make the API request
        response = requests.get(base_url, params=params)
        response.raise_for_status()  # Raise an exception for bad status codes
        
        return response.json()
    
    except requests.exceptions.RequestException as e:
        return {"error": str(e)}

def detect_id_type(id_string: str) -> str:
    """
    Detect the type of the given ID or title.
    
    Args:
    id_string (str): The ID or title to detect.
    
    Returns:
    str: The detected type ('pmid', 'pmcid', 'doi', 'openalex', 'semantic_scholar', 'potential_title', or 'unknown').
    """
    if re.match(r'^\d{1,8}$', id_string):
        return 'pmid'
    elif re.match(r'^PMC\d+$', id_string):
        return 'pmcid'
    elif re.match(r'^10\.\d{4,9}/[-._;()/:A-Z0-9]+$', id_string, re.I): # Detect DOIs, case-insensitive
        return 'doi'
    elif re.match(r'^[WAIC]\d{2,}$', id_string):
        return 'openalex'
    elif re.match(r'^[0-9a-f]{40}$', id_string): 
        return 'semantic_scholar'
    elif re.match(r'^[A-Z][\w\s:,\-()]{10,150}[.?!]?$', id_string, re.I): # A simple heuristic for detecting titles
        return 'potential_title'
    else:
        return 'unknown'


@json_schema_decorator
def id_converter_tool(ids: List[str]) -> str:
    """
    For any article(s) in PubMed Central, find all the corresponding PubMed IDs (PMIDs), digital object identifiers (DOIs), and manuscript IDs (MIDs).
    Use this tool to convert a list of IDs, such as PMIDs, PMCIDs, or DOIs, and find the corresponding IDs for the same articles.
    
    Args:
    ids (str): A string with a comma-separated list of IDs to convert. Must be PMIDs, PMCIDs, or DOIs. The maximum number of IDs per request is 200.
    
    Returns:
    str: A JSON-formatted string containing the conversion results and the detected ID types.
    """

    if isinstance(ids, str):
        ids = ids.replace("https://doi.org/", "").replace("doi.org/", "") # Remove DOI URL prefixes, if present
    
    # Convert the input string to a list of IDs
    try:
        ids = ids.split(",")
    except AttributeError:
        return json.dumps({"error": "Input must be a string of comma-separated IDs"})

    if len(ids) > 200:
        return json.dumps({"error": "Input IDs must be no more than 200"})
    
    # Detect ID types
    id_types = [detect_id_type(id_string) for id_string in ids]
    
    # Convert IDs
    conversion_result = convert_id(ids)
    
    # Check if conversion_result is already a dictionary (error case)
    if isinstance(conversion_result, dict):
        parsed_result = conversion_result
    else:
        # Parse the JSON string result
        try:
            parsed_result = json.loads(conversion_result)
        except json.JSONDecodeError:
            return json.dumps({"error": "Failed to parse conversion result"})
    
    # Prepare the result
    result = {
        "conversion_result": parsed_result,
        "detected_id_types": dict(zip(ids, id_types))
    }
    
    return json.dumps(result, indent=2)

# %% ../nbs/01_core.ipynb 98
@json_schema_decorator
def query_openalex_api(query_param: str) -> str:
    """
    Retrieve metadata for a given article from OpenAlex, a comprehensive open-access catalog of global research papers.
    Use this tool to search the OpenAlex API by using the article title, the PubMed ID (PMID), the PubMed Central ID (PMCID) or the digital object identifier (DOI) of an article as the query parameter. 
    This tool returns the following metadata:
    - the OpenAlex ID
    - the digital object identifier (DOI) URL
    - Citation count
    - The open access status
    - URL to the open-access location for the work
    - Publication year
    - A URL to a website listing works that have cite the article
    - The type of the article
    Use this tool only if an article title, PubMed ID or DOI is provided by the user or was extracted from a local PDF file and is present in the conversation history.
    
    Args:
        query_param (str): The article title, the PubMed ID (PMID), the PubMed Central ID (PMCID) or the digital object identifier (DOI) of the article to retrieve metadata for. May be provided by the user or extracted from a local PDF file and present in the conversation history.

    Returns:
        str: A JSON-formatted string including the search results from the OpenAlex database. If no results are found or the API query fails, an appropriate message is returned.
    """

    if 'EMAIL' not in globals():
        try:
            EMAIL = os.getenv("EMAIL")
        except KeyError:
            EMAIL = ""

    # Validate the input
    if query_param is None or query_param == "":
        return json.dumps({"error": "The query parameter must be a non-empty string."})
    elif not isinstance(query_param, str):
        query_param = str(query_param)
    
    if query_param.startswith("https://doi.org/"):
        query_param = query_param.replace("https://doi.org/", "")
    elif query_param.startswith("doi.org/"):
        query_param = query_param.replace("doi.org/", "")

    # Constants
    base_url = "https://api.openalex.org/" # Define the base URL for the OpenAlex API
    
    # Initialize variables
    filter = ""

    if detect_id_type(query_param) == "potential_title":
        url = f"{base_url}works?"
        filter = f"title.search:{query_param}"
    elif detect_id_type(query_param) == "pmid":
        url = f"{base_url}works/pmid:{query_param}"
    elif detect_id_type(query_param) == "doi":
        url = f"{base_url}works/https://doi.org/{query_param}"
    elif detect_id_type(query_param) == "pmcid":
        url = f"{base_url}works/pmcid:{query_param}"
    else:
        return json.dumps({"error": "The query parameter must be a valid title, PMID, PMCID, or DOI."})
        
    # Set the query parameters
    params = {
        "mailto": EMAIL,
        "page": 1,
        "per-page": 5,
        "select": "id,doi,title,publication_year,cited_by_count,cited_by_api_url,open_access,type,type_crossref",
    }
    if filter:
        params["filter"] = filter # Add the filter parameter for title search

    response = requests.get(url, params=params)

    if response.status_code != 200:
        return {"error": f"Failed to query OpenAlex API. Status code: {response.status_code}"}

    raw_search_results = response.json()

    if filter:
        number_of_search_matches = raw_search_results['meta']['count']
        if len(raw_search_results['results']) == 0:
            return "No results found for the provided title."
        elif number_of_search_matches > 5:
            return "Error: The search results are more than 5. Please provide a correct title."
        else:
            formatted_results = []
            for result in raw_search_results['results']:
                formatted_results.append(result)
            return json.dumps(formatted_results, indent=2)

    
    if len(raw_search_results) == 0:
        return {"search result": "None"}
    else:
        return json.dumps(raw_search_results, indent=2)

# %% ../nbs/01_core.ipynb 113
import time
import re

# %% ../nbs/01_core.ipynb 115
@json_schema_decorator
def query_semantic_scholar_api(query_param: str) -> str:
    """
    Retrieve metadata for a given article from the Semantic Scholar Academic Graph (S2AG), a large knowledge graph of scientific literature that combines data from multiple sources.
    Use this tool to query the Semantic Scholar Graph API by using either the article title, the PubMed ID, or the digital object identifier (DOI) to retrieve the following metadata:
    - the title
    - the publication year
    - the abstract
    - a tldr (too long, didn't read) summary
    - the authors of the article
    - the URL to the open-access PDF version of the article, if available
    - the journal name
    - a url to the article on the Semantic Scholar website
    Use this tool only if an article title, PubMed ID or DOI is provided by the user or was extracted from a local PDF file and is present in the conversation history.
    
    Args:
        query_param (str): The article title, the PubMed ID, or the digital object identifier of the article to retrieve metadata for. May be provided by the user or extracted from a local PDF file and present in the conversation history. Do not include the 'https://doi.org/' prefix for DOIs, or keys such as 'DOI', 'PMCID' or 'PMID'. The tool will automatically detect the type of identifier provided.

    Returns:
        str: A JSON-formatted string including the search results from the Semantic Scholar database. If no results are found or the API query fails, an appropriate message is returned.
    """

    # Try to get the API key from the environment variables, if available, and define the headers
    try:
        SEMANTIC_SCHOLAR_API_KEY
    except NameError:
        SEMANTIC_SCHOLAR_API_KEY = os.environ.get("SEMANTIC_SCHOLAR_API_KEY")
    headers = {'x-api-key': SEMANTIC_SCHOLAR_API_KEY}

    # Validate the input
    if query_param is None or query_param == "":
        return json.dumps({"error": "The query parameter must be a non-empty string."})
    elif not isinstance(query_param, str):
        query_param = str(query_param)
    
    # Clean the query parameter
    query_param = query_param.strip().lower() # Convert to lowercase and remove leading/trailing whitespace
    if query_param.startswith("https://doi.org/"):
        query_param = query_param.replace("https://doi.org/", "")
    elif query_param.startswith("doi.org/"):
        query_param = query_param.replace("doi.org/", "")
    elif query_param.startswith("doi"):
        query_param = query_param.replace("doi", "")
    elif query_param.startswith("pmid"):
        query_param = query_param.replace("pmid", "")
    elif query_param.startswith("pmcid"):
        query_param = query_param.replace("pmcid", "")
    if detect_id_type(query_param) != "potential_title":
        query_param = query_param.replace(":", "") # Remove any colons from the query parameter if it is not a title
        query_param = query_param.replace(" ", "") # Remove spaces from the query parameter if it is not a title

    # Constants
    academicgraph_base_url = "https://api.semanticscholar.org/graph/v1"
    fields = "title,year,authors,tldr,abstract,citationCount,openAccessPdf,journal,url" # The values to retrieve from the API


    # Construct the URL based on the identifier type
    if detect_id_type(query_param) == "potential_title":
        url = f"{academicgraph_base_url}/paper/search/match?query={query_param}&fields={fields}"
    elif detect_id_type(query_param) == "pmid":
        url = f"{academicgraph_base_url}/paper/PMID:{query_param}?&fields={fields}"
    elif detect_id_type(query_param) == "doi":
        url = f"{academicgraph_base_url}/paper/DOI:{query_param}?&fields={fields}"
    else:
        return json.dumps({"error": "The query parameter must be a valid title, PMID, or DOI."})
    
    response = requests.get(url, headers=headers) # Send the API request
    # print(response.url) # Uncomment for debugging
    time.sleep(2) # sleep for 2 seconds to avoid rate limiting issues
    
    # Check response status
    if response.status_code == 200:
        return json.dumps(response.json(), indent=2)
    else:
        return f"Error: Failed to query Semantic Scholar API. Status code: {response.status_code}"

# %% ../nbs/01_core.ipynb 122
@json_schema_decorator
def respond_to_generic_queries() -> str:
    """
    A function to respond to generic questions or queries from the user. Use this tool if no better tool is available.

    This tool does not take any arguments.

    Returns:
        str: A response to a generic question.
    """

    return "There is no specific tool available to respond this query from the user. State your capabilities based the system message or provide a response based on the conversation history."

# %% ../nbs/01_core.ipynb 126
from typing import Generator

# %% ../nbs/01_core.ipynb 129
import ollama
from typing import Dict, Any, List

# %% ../nbs/01_core.ipynb 130
import sys
import io
from pathlib import Path

class Assistant:
    def __init__(self,
        status: dict = {}, # The status of the assistant
        sys_message: str or None = None, # The system message for the assistant; if not provided, a default message is used
        model: str = "llama3.1:latest", # The model to use for the assistant
        tools: Dict[str, Any] = { # The tools available to the assistant
           "get_file_names": get_file_names,
           "extract_text_from_pdf": extract_text_from_pdf,
           "get_titles_and_first_authors": get_titles_and_first_authors,
           "summarize_local_document": summarize_local_document,
           "describe_python_code": describe_python_code,
           "id_converter_tool": id_converter_tool,
           "query_openalex_api": query_openalex_api,
           "query_semantic_scholar_api": query_semantic_scholar_api,
           "respond_to_generic_queries": respond_to_generic_queries,
        },
        add_tools: Dict[str, Any] = {}, # Optional argument to add additional tools to the assistant, when initializing
        authentication: Optional[Dict[str, str]] = None, # Authentication credentials for API calls to external services
        dir_path: str = "../data", # The directory path to which the assistant has access on the local computer
        messages: List[Dict[str, str]] = []): # The conversation history
        
        # ANSI escape codes, used for output formatting
        GREY = "\033[90m"
        BLUE = "\033[94m"
        ORANGE = "\033[33m"
        RED = "\033[91m"
        RESET = "\033[0m"

        self.model = model

        self.sys_message = sys_message
        self.tools = tools
        self.tools.update(add_tools) # Add additional tools to the assistant, if provided
        if self.tools:
            self.tools["describe_tools"] = self.describe_tools # Add the describe_tools function to the tools list for the assistant, if the tools list is not empty
        self.authentication = authentication or {}
        self.dir_path = Path(dir_path).resolve()
        self.messages = messages
        self.status = self.get_status()
        if self.status["status_code"] not in [2, 3]:
            print(f"{self.status['detail']}")
            self.server_available = False
            return None
        else:
            if self.status["status_code"] == 3:
                print(f"{self.status['detail']}")
            self.server_available = True

        # Set global variables
        global DIR_PATH
        DIR_PATH = self.dir_path
        global MODEL
        MODEL = self.model
        # TODO: Consider allowing the user to set different models for different tasks and tools
        # e.g. a model such as llama 3.1 for function calls, command-r-plus for summarization, aya for translation, etc.
        
        # Load the API keys from the environment variables or the authentication dictionary
        self.SEMANTIC_SCHOLAR_API_KEY = self.authentication.get("SEMANTIC_SCHOLAR_API_KEY")
        self.EMAIL = self.authentication.get("EMAIL")

        if not self.SEMANTIC_SCHOLAR_API_KEY:
            self.SEMANTIC_SCHOLAR_API_KEY = os.environ.get("SEMANTIC_SCHOLAR_API_KEY")
            if self.SEMANTIC_SCHOLAR_API_KEY:
                print(f"{GREY}Loaded Semantic Scholar API key from the environment variables.{RESET}")
        if not self.EMAIL:
            self.EMAIL = os.environ.get("EMAIL")
            if self.EMAIL:
                print(f"{GREY}Loaded email address from the environment variables.{RESET}")

        # Generate the default directory for storing data files if it does not exist
        if not os.path.exists(DIR_PATH):
            os.mkdir(DIR_PATH)
            print(f"{GREY}Created directory {DIR_PATH} for storing data files.{RESET}\n")
        else: 
            print(f"{GREY}A local directory {DIR_PATH} already exists for storing data files. No of files: {len(os.listdir(DIR_PATH))}{RESET}\n")

        # Set the default system message if not provided
        if not self.sys_message:
            self.sys_message ="""You are an AI assistant specialized in analyzing research articles.
        Your role is to provide concise, human-readable responses based on information from tools and conversation history.

        Key instructions:
        1. Use provided tools to gather information before answering.
        2. Interpret tool results and provide clear, concise answers in natural language.
        3. If you can't answer with available tools, state this clearly.
        4. Don't provide information if tool content is empty.
        5. Never include raw JSON, tool outputs, or formatting tags in responses.
        6. Format responses as plain text for direct human communication.
        7. Use clear formatting (e.g., numbered or bulleted lists) when appropriate.
        8. Provide article details (e.g., DOI, citation count) in a conversational manner.

        Act as a knowledgeable research assistant, offering clear and helpful information based on available tools and data.
        """

        # Check if tools are provided
        if len(self.tools) == 0:
            print(f"\033[91mNo tools provided! Please add tools to the assistant.\033[0m")

    def __str__(self):
        return f"Assistant, powered by {self.model.split(':')[0]}"

    def __repr__(self):
        return self.__str__()

    def get_status(self):
        """Get the status of the assistant initialization."""
        status = {}
        try:
            response = requests.get("http://localhost:11434")
            if response.status_code == 200:
                downloaded_models = [model.get("model").split(":")[0] for model in ollama.list()["models"]]
                tool_supporting_models = ["llama3.2", "llama3.1", "qwen2.5", "mistral-nemo", "nemotron-mini", "command-r", "command-r-plus", "firefunction-v2"]
                if not any(model in downloaded_models for model in tool_supporting_models):
                    status["status_code"] = 3
                    status["detail"] = "Warning: Unknown model. Did you download a model that supports tool calling? Visit https://ollama.com/search?c=tools for details."
                elif len(downloaded_models) == 0:
                    status["status_code"] = 4
                    status["detail"] = "Error: No local models found. Visit https://ollama.com/models for available models and instructions on how to download them."
                elif not hasattr(self, "model") or self.model.split(":")[0] not in downloaded_models:
                    status["status_code"] = 5
                    status["detail"] = "Error: Selected model not downloaded. Visit https://ollama.com/models for available models and instructions on how to download them."
                else:
                    status["status_code"] = 2
                    status["detail"] = "Assistant initialized successfully."
            else:
                status["status_code"] = 6
                status["detail"] = "Error: Could not connect to local host at port 11434. Is the Ollama app running? Aborting initialization of the assistant."
        except Exception as e:
            if "Connection refused" in str(e):
                status["status_code"] = 7
                status["detail"] = "Error: Could not connect to a running Ollama instance. Is the Ollama app installed and the server running? Aborting initialization of the assistant."
            else:
                status["status_code"] = 8
                status["detail"] = f"Error: {e}"
        return status

    def list_tools(self):
        "List the available tools in the assistant."
        for tool in self.tools.keys():
            print(tool)

    def get_tools_schema(self):
        "Return the JSON schema for the available tools."
        return [func.json_schema for func in self.tools.values()]

    @json_schema_decorator
    def describe_tools(self) -> str:
        """Use this tool when asked about the assistant's available tools and capabilities.

        Returns:
            str: A string with the descriptions of the available tools.
        """
        return f"Available tools are: {self.get_tools_schema()}\n State your capabilities based the available tools in a conversational manner."

    def chat(self, prompt: str, show_progress: bool = False, stream_response: bool = True, redirect_output: bool = False):
        """
        Start a conversation with the AI assistant.

        Args:
            prompt (str): The user's prompt or question.
            show_progress (bool): Whether to show the step-by-step progress of the fuction calls, including the tool calls and tool outputs. Default is False.
            stream_response (bool): Whether to stream the final response from the LLM. Default is True. Automatically set to True if redirect_output is True.
            redirect_output (bool): Whether to redirect the output to be compatible with st.write_stream. Default is False.

        Returns:
            str: The AI assistant's response.
        """
        # Abort the function if the server is not available
        if not self.server_available:
            print(self.get_status()["detail"])
            return None

        # At the start of the conversation, if no messages are provided, add the system message and user prompt
        if not self.messages:
            self.messages = [
                {'role': "system", 'content': self.sys_message},
                {'role': 'user', 'content': prompt},
            ]
        else:
            self.messages.append({'role': 'user', 'content': prompt})

        # Generate JSON schemas for the available tools
        tools_schema = self.get_tools_schema()

        # Make a request to the LLM to select a tool
        if show_progress: print("Selecting tools...\n")
        response = ollama.chat(
            model=self.model,
            messages=self.messages,
            tools=tools_schema,
            stream=False, # Set to False to avoid streaming the tool calls
        )

        # Add the model's response to the conversation history
        if response.get('message', {}).get('tool_calls'):
            if show_progress: print(response['message']['tool_calls']) # Uncomment for debugging
            self.messages.append(
                {'role': 'assistant', 'tool_calls': response['message']['tool_calls']}
                )
        else:
            # print("LLM response (not added to the conversation history):", response['message']['content']) # Uncomment for debugging
            self.messages.append(
                {'role': 'assistant', 'tool_calls': []}
                )
            print(f"\033[91mNo tool calls found in the response. Adding an empty tool_calls list to the conversation history. Aborting...\033[0m\n")
            return None # Abort the function if no tool calls are found in the response. Goal is to force the assistant to use a tool. We will generate a tool for generic responses.

        # Call the function if a tool is selected
        for tool in response['message']['tool_calls']:
            # print("Arguments:", tool['function']['arguments'])
            function_to_call = self.tools[tool['function']['name']]
            if show_progress: print(f"Calling {tool['function']['name']}() with arguments {tool['function']['arguments']}...\n")
            args = tool['function']['arguments']

            try:
                function_response = function_to_call(**args)
                # print(f"Function response type: {type(function_response)}\n") # Uncomment for debugging
                # print(f"Function response: {function_response}\n") # Uncomment for debugging
                function_response = str(function_response) if function_response is not None else ""
                # assert isinstance(function_response, str), "Function response must be a string."
            except Exception as e:
                function_response = f"Error: {e}"

                if show_progress: print(f"Function response:\n{function_response}\n")
            # Add the fucntion response to the conversation history
            self.messages.append( 
                {
                    'role': 'tool',
                    'content': function_response,
                }
            ) 

        if redirect_output: # If the output is to be redirected...
            stream_response = True # always use streaming for compatibility with st.write_stream

        # Make a second request to the LLM with the tool output to generate a final response
        if show_progress: print("Generating final response...")
        response = ollama.chat(
            model=self.model,
            format="", # Set to empty string to avoid JSON formatting; If JSON formatting is needed, set to "json"
            messages=self.messages,
            stream=stream_response, # Set to True , response will be a generator
        )
        
        # Advanced parameters (optional):
        # keep_alive: controls how long the model will stay loaded into memory following the request (default: 5m)
        # options: additional model parameters listed in the documentation for the Modelfile such as temperature
        # for more information, visi:
        # https://github.com/ollama/ollama/blob/main/docs/api.md
        # https://github.com/ollama/ollama/blob/main/docs/modelfile.md#valid-parameters-and-values
        
        if redirect_output:
            # Return a generator for st.write_stream
            def response_generator():
                full_content = ""
                for chunk in response:
                    content = chunk['message']['content']
                    full_content += content
                    yield content
                self.messages.append({'role': 'assistant', 'content': full_content})
            return response_generator()

        else:
            if isinstance(response, dict):
                content = response['message']['content']
                self.messages.append({'role': 'assistant', 'content': content})
                print(content)
                return content  # Return the content directly, not as a generator
                
            elif hasattr(response, '__iter__'):  # Check if it's iterable (for streaming)
                full_content = ""
                for chunk in response:
                    content = chunk['message']['content']
                    full_content += content
                    print(content, end='', flush=True)
                self.messages.append({'role': 'assistant', 'content': full_content})
                return full_content  # Return the full content, not as a generator
    
    def show_conversion_history(self, show_function_calls: bool = False):
        """Display the conversation history.
        
        Args:
            show_function_calls (bool): Whether to show function calls and returns in the conversation history. Default is False.

        Returns:
            None
        """

        # ANSI escape code for blue and red text
        BLUE = "\033[94m"
        BOLD = "\033[1m"
        GREY = "\033[90m"
        RESET = "\033[0m"

        for message in self.messages:
            if message['role'] != 'system':
                if message['role'] == 'user':
                    print(f"{BOLD}User:{RESET} {message['content']}\n")
                elif message['role'] == 'assistant':
                    if 'content' in message and message['content']:
                        print(f"{BOLD}{BLUE}Assistant response:{RESET} {BLUE}{message['content']}{RESET}\n")
                    if 'tool_calls' in message and message['tool_calls'] and show_function_calls:
                        print(f"{BOLD}{BLUE}Assistant function calls:{RESET} ", end='')
                        for tool in message['tool_calls']:
                            print(f"{BLUE}{tool['function']['name']}() with arguments {tool['function']['arguments']}{RESET}\n")
                elif message['role'] == 'tool' and show_function_calls:
                    # convert str to list
                    if isinstance(message['content'], str):
                        message['content'] = [message['content']]
                    for fn_return in message['content']:
                        print(f"{BOLD}{GREY}Function return:{RESET} {GREY}{fn_return}{RESET}\n")
    
    def clear_conversion_history(self):
        """Clear the conversation history."""
        self.messages = [{'role': "system", 'content': self.sys_message},]
    
    def pprint_tools(self):
        """Pretty-print the descriptions of the available tools."""
        for tool in self.get_tools_schema():   
            print(f"""* Tool name: {tool.get("function", {}).get("name", "No name available.")}
        Description: {tool.get("function", {}).get("description", "No description available.")}
            """)
        return None

# %% ../nbs/01_core.ipynb 143
def add_to_class(Class: type):
    """Register functions as methods in a class that has already been defined."""
    def wrapper(obj):
        setattr(Class, obj.__name__, obj)
    return wrapper
