Upload 22 files

.gitattributes

@@ -33,6 +33,8 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+img/11111111.png filter=lfs diff=lfs merge=lfs -text
+img/2222222.png filter=lfs diff=lfs merge=lfs -text
 *.nc filter=lfs diff=lfs merge=lfs -text
 *.sqlite3 filter=lfs diff=lfs merge=lfs -text
 *.shp filter=lfs diff=lfs merge=lfs -text

src/agents.py

@@ -0,0 +1,1097 @@
+# src/agents.py
+import base64
+import os
+import logging
+import functools
+import subprocess
+import sys
+from io import StringIO
+from typing import List, Annotated, Sequence, TypedDict
+import operator
+from langchain.agents import AgentExecutor, create_openai_tools_agent
+from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
+from langchain_core.output_parsers.openai_functions import JsonOutputFunctionsParser
+from langchain_core.messages import BaseMessage, HumanMessage
+from langchain_experimental.tools import PythonREPLTool
+from langchain_openai import ChatOpenAI, OpenAI
+from langchain_core.messages import AIMessage
+from langchain_experimental.agents.agent_toolkits.pandas.base import create_pandas_dataframe_agent
+from pydantic import BaseModel, Field, PrivateAttr
+from langchain_core.tools import StructuredTool
+from langchain.agents.format_scratchpad.openai_tools import (
+    format_to_openai_tool_messages,
+)
+from langchain.agents.output_parsers.openai_tools import OpenAIToolsAgentOutputParser
+from langgraph.graph import StateGraph, END
+from langchain.agents.agent_types import AgentType
+
+
+from langchain_openai import OpenAIEmbeddings
+from langchain_community.vectorstores import Chroma
+
+from typing import Any
+import streamlit as st
+import pandas as pd
+import matplotlib.pyplot as plt
+from openai import OpenAI
+
+
+#sys.path
+# Get the absolute path of the current file (agents.py)
+current_dir = os.path.dirname(os.path.abspath(__file__))
+
+# Get the parent directory
+parent_dir = os.path.abspath(os.path.join(current_dir, '..'))
+
+# Add the parent directory to sys.path
+if parent_dir not in sys.path:
+    sys.path.insert(0, parent_dir)
+
+
+# Import custom modules
+from .search.search_pg_default import pg_search_default
+from .search.publication_qa_tool import answer_publication_questions, PublicationQAArgs
+from .plotting_tools.hard_agent import plot_master_track_map
+from .plotting_tools.oceanographer_tools import plot_ts_diagram
+from .prompts import Prompts
+from .utils import generate_unique_image_path
+from .config import API_KEY
+
+
+# 1. Search Agent and Tools
+class CustomPythonREPLTool(PythonREPLTool):
+    _datasets: dict = PrivateAttr()
+
+    def __init__(self, datasets, **kwargs):
+        """
+        Custom Python REPL tool that injects dataset variables and logs plot generation.
+        :param datasets: Dictionary { "dataset_1": <DataFrame>, "dataset_2": <DataFrame>, ... }
+        """
+        super().__init__(**kwargs)
+        self._datasets = datasets
+
+    def _run(self, query: str, **kwargs) -> Any:
+        """
+        Execute the user-provided Python code in a local context containing:
+          - st (Streamlit)
+          - plt (Matplotlib Pyplot)
+          - pd (Pandas)
+          - All loaded dataset variables (self._datasets)
+          - A dynamically generated plot_path
+
+        If a figure is saved to plot_path, a "plot_generated" event will be logged in session_state["execution_history"].
+        """
+        import streamlit as st
+        import matplotlib.pyplot as plt
+        import pandas as pd
+        import logging
+        from io import StringIO
+        from src.utils import log_history_event, generate_unique_image_path
+
+        # Prepare local context with necessary packages
+        local_context = {"st": st, "plt": plt, "pd": pd}
+
+        # Inject the user’s datasets under the specified variable names (e.g. dataset_1, dataset_2, etc.)
+        local_context.update(self._datasets)
+
+        # Generate a unique file path for the plot (plot_path)
+        plot_path = generate_unique_image_path()
+        local_context['plot_path'] = plot_path
+
+        # Redirect stdout so we can capture any output from exec(...)
+        old_stdout = sys.stdout
+        redirected_output = StringIO()
+        sys.stdout = redirected_output
+
+        try:
+            # Execute user code
+            exec(query, local_context)
+            output = redirected_output.getvalue()
+
+        except ModuleNotFoundError as e:
+            missing_module = e.name
+            logging.warning(f"Module '{missing_module}' not found during code execution.")
+            return {
+                "error": "ModuleNotFoundError",
+                "missing_module": missing_module,
+                "message": f"The Python module '{missing_module}' is not installed."
+            }
+        except Exception as e:
+            logging.error(f"Error during code execution: {e}")
+            return {
+                "error": "ExecutionError",
+                "message": str(e)
+            }
+        finally:
+            # Restore stdout
+            sys.stdout = old_stdout
+
+        # Check if a plot was actually saved to plot_path
+        plot_generated = False
+        if os.path.exists(plot_path):
+            st.session_state.saved_plot_path = plot_path
+            st.session_state.plot_image = plot_path
+            st.session_state.new_plot_path = plot_path
+            plot_generated = True
+
+        if plot_generated:
+            status_message = f"Plot generated = True. Saved at: {plot_path}"
+            logging.info(status_message)
+            st.session_state.plot_generated_status = status_message
+            
+            from src.utils import log_history_event
+            log_history_event(
+                st.session_state,
+                "plot_generated",
+                {
+                    "plot_path": plot_path.replace("sandbox:", ""),  # Remove sandbox prefix
+                    "agent": "VisualizationAgent",
+                    "description": "Python_REPL Generated Plot",
+                    "content": query  # Store the actual code used
+                }
+            )
+
+        return {
+            "result": f"Execution completed. Plot saved at: {plot_path if plot_generated else 'No plot generated'}",
+            "output": output,
+            "plot_images": [plot_path] if plot_generated else []
+        }
+
+
+
+def search_pg_datasets_tool(query):
+    global prompt_search
+
+    datasets_info = pg_search_default(query)
+
+    logging.debug("Datasets info: %s", datasets_info)
+
+    if not datasets_info.empty:
+        st.session_state.datasets_info = datasets_info
+        st.session_state.messages_search.append({
+            "role": "assistant", 
+            "content": f"**Search query:** {query}"
+        })
+        # Pass the table as JSON (you can use orient="split" or the default, as long as your UI can parse it)
+        st.session_state.messages_search.append({
+            "role": "assistant", 
+            "content": "**Datasets Information:**", 
+            "table": datasets_info.to_json(orient="split")
+        })
+
+        # Optionally, build a detailed description string for the prompt:
+        datasets_description = ""
+        for i, row in datasets_info.iterrows():
+            datasets_description += (
+                f"Dataset {i + 1}:\n"
+                f"Name: {row['Name']}\n"
+                f"Description: {row['Short Description']}\n"
+                f"Parameters: {row['Parameters']}\n\n"
+            )
+
+        prompt_search = (
+            f"The user has provided the following query: {query}\n"
+            f"Available datasets:\n{datasets_description}\n"
+            "Please identify the top two datasets that best match the user's query and explain why they are the most relevant. "
+            "Do not suggest datasets without values in the Parameters field, because they cannot be directly downloaded.\n"
+            "Respond using the following schema:\n"
+            "{dataset name}\n{reason why relevant}\n{propose some short analysis and further questions to answer}"
+        )
+
+    return datasets_info, prompt_search
+
+
+def create_search_agent(datasets_info=None):
+    model_name = st.session_state.get("model_name", "gpt-3.5-turbo")
+    if model_name == "o3-mini":
+        llm = ChatOpenAI(api_key=API_KEY, model_name=model_name)
+    else:
+        llm = ChatOpenAI(api_key=API_KEY, model_name=model_name)
+
+    # Generate dataset description string
+    datasets_description = ""
+    if datasets_info is not None:
+        for i, row in datasets_info.iterrows():
+            datasets_description += f"Dataset {i + 1}:\nName: {row['Name']}\nDOI: {row['DOI']}\nDescription: {row['Short Description']}\nParameters: {row['Parameters']}\n\n"
+
+    prompt = ChatPromptTemplate.from_messages(
+        [
+            ("system",
+             f"You are a powerful assistant primarily designed to search and retrieve datasets from PANGAEA. Your main goal is to help users find relevant datasets using the search_pg_datasets tool. When a user asks about datasets, always use this tool first to provide the most up-to-date and accurate information.\n\n"
+             #f"Here are some datasets returned from the search:\n{datasets_description}"
+             "In addition to dataset searches, you have a secondary capability to answer questions about publications related to specific datasets (or in other words what was published based on this dataset). If a user explicitly asks about publications or research findings based on a particular dataset, you can use the answer_publication_questions tool. For example, you can handle queries like 'What was published based on this dataset?' or 'What were the main conclusions from the research using this dataset?'\n\n"
+             "Remember:\n"
+             "1. Prioritize dataset searches using the search_pg_datasets tool. Make sure that the query you pass to the tool is rephrased so that elastic search gives the best match. Also try not to include words like 'search' and etc. in the search query.\n"
+             "2. Only use the answer_publication_questions tool when the user specifically asks about publications or research findings related to a dataset they've already identified. Please make sure that you correctly pass the doi to the tool. It should be doi retrieved after the search (user will point out which dataset it interested in). DO NOT GENERATE DOI ON THIS STEP OUT OF YOUR MIND! JUST TAKE WHAT WAS GIVEN WITH SYSTEM PROMPT.\n"
+             "3. If needed, ask the user to clarify which dataset they're referring to before using the publication tool.\n\n"
+             "Strive to provide accurate, helpful, and concise responses to user queries."
+             ),
+            ("user", "{input}"),
+            MessagesPlaceholder(variable_name="chat_history"),
+            MessagesPlaceholder(variable_name="agent_scratchpad"),
+        ]
+    )
+
+    search_tool = StructuredTool.from_function(
+        func=search_pg_datasets_tool,
+        name="search_pg_datasets",
+        description="List datasets from PANGAEA based on a query"
+    )
+
+    publication_qa_tool = StructuredTool.from_function(
+        func=answer_publication_questions,
+        name="answer_publication_questions",
+        description="A tool to answer questions about articles published from this dataset. This will be a journal article for which you should provide the tool with an already structured question about what the user wants. The input should be the DOI of the dataset (e.g. 'https://doi.org/10.1594/PANGAEA.xxxxxx') and the question. The question should be reworded to specifically send it to RAG. E.g. the hypothetical user's question 'Are there any related articles to the first dataset? If so what these articles are about?' will be re-worded for this tool as 'What is this article is about?'",
+        args_schema=PublicationQAArgs
+    )
+
+    tools = [search_tool, publication_qa_tool]
+
+    llm_with_tools = llm.bind_tools(tools)
+
+    agent = (
+            {
+                "input": lambda x: x["input"],
+                "chat_history": lambda x: x.get("chat_history", []),
+                "agent_scratchpad": lambda x: format_to_openai_tool_messages(x["intermediate_steps"]),
+            }
+            | prompt
+            | llm_with_tools
+            | OpenAIToolsAgentOutputParser()
+    )
+
+    return AgentExecutor(agent=agent, tools=tools, verbose=True, max_iterations=5)
+
+# 2. Visualization and Oceanography Tools
+
+class PlotMasterTrackMapArgs(BaseModel):
+    dataset_var: str = Field(description="The variable name of the dataset to use (e.g., 'dataset_1', 'dataset_2').")
+    main_title: str = Field(description="The main title for the plot.")
+    lat_col: str = Field(description="Name of the latitude column.")
+    lon_col: str = Field(description="Name of the longitude column.")
+    date_col: str = Field(description="Name of the date/time column.")
+
+class TSPlotToolArgs(BaseModel):
+    dataset_var: str = Field(description="The variable name of the dataset to use (e.g., 'dataset_1', 'dataset_2').")
+    main_title: str = Field(description="The main title for the plot.")
+    temperature_col: str = Field(description="Name of the temperature column.")
+    salinity_col: str = Field(description="Name of the salinity column.")
+
+# 3. Agent Creation Functions
+def create_pandas_agent(user_query, datasets_info):
+    if st.session_state.model_name == "o3-mini":
+        llm = ChatOpenAI(api_key=API_KEY, model_name=st.session_state.model_name)
+    else:
+        llm = ChatOpenAI(api_key=API_KEY, model_name=st.session_state.model_name)
+
+    # Assign unique variable names to each dataframe and collect dataframes
+    dataset_variables = []
+    dataframes = []
+    datasets_text = ""  # Initialize datasets_text
+    for i, info in enumerate(datasets_info, 1):  # Start enumeration at 1
+        var_name = f"df{i}"  # Consistently name as df1, df2, etc.
+        dataframes.append(info['dataset'])  # Collect dataframes into a list
+        dataset_variables.append(var_name)
+        # Build datasets_text
+        datasets_text += (
+            f"Dataset {i}:\n"  # Adjust index to match variable naming
+            f"Variable Name: {var_name}\n"
+            f"Name: {info['name']}\n"
+            f"Description: {info['description']}\n"
+            f"Head of DataFrame (use it only as an example):\n"
+            f"{info['df_head']}\n\n"
+        )
+
+
+    # Create a custom system prompt that includes information about each dataframe
+    system_prompt = Prompts.generate_pandas_agent_system_prompt(user_query, datasets_text, dataset_variables)
+
+    # Create a ChatPromptTemplate with the system prompt
+    chat_prompt = ChatPromptTemplate.from_messages(
+        [
+            ("system", system_prompt),
+            ("user", "{input}"),
+            MessagesPlaceholder(variable_name="chat_history"),
+            MessagesPlaceholder(variable_name="agent_scratchpad"),
+        ]
+    )
+
+    # Create the pandas dataframe agent with the list of dataframes and the chat prompt
+    agent_pandas = create_pandas_dataframe_agent(
+        llm=llm,
+        df=dataframes,  # Pass the list of dataframes
+        agent_type=AgentType.OPENAI_FUNCTIONS,
+        verbose=True,
+        return_intermediate_steps=True,
+        max_iterations=5,
+        early_stopping_method="generate",
+        handle_parsing_errors=True,
+        #prefix=system_prompt,
+        suffix=system_prompt,
+        allow_dangerous_code=True,
+        chat_prompt=chat_prompt
+    )
+
+    return agent_pandas
+
+
+# Define the function to encode the image
+def encode_image(image_path):
+    with open(image_path, "rb") as image_file:
+        return base64.b64encode(image_file.read()).decode('utf-8')
+
+def reflect_on_image(image_path: str) -> str:
+    if not os.path.exists(image_path):
+        return f"Error: The file {image_path} does not exist."
+
+    base64_image = encode_image(image_path)
+
+    prompt = """You are a professional reviewer of scientific images. Your task is to provide review and pass it back to the visual creator agent, so that it could improve it. At each step provide idea for improvements (only if necessary). Be sure to be critical and provide a source for improvement. Conduct a quality check of the provided image using the following criteria:
+
+1. Axis and Font Quality: Evaluate the visibility of axes and appropriateness of font size and style. Are the axes clearly visible and labeled? Is the font legible and suitable for the image size?
+2. Label Clarity: Assess if labels are well-positioned and not overlapping. Are all labels clearly readable and properly placed?
+3. Color Scheme: Analyze the color choices. Is the color scheme appropriate for the data presented? Are the colors distinguishable and not causing visual confusion?
+4. Data Representation: Evaluate how well the data is represented. Are data points clearly visible? Is the chosen chart or graph type appropriate for the data?
+5. Legend and Scale: Check the presence and clarity of legends and scales. Are they present when necessary and easy to understand?
+6. Overall Layout: Assess the overall layout and use of space. Is the image well-balanced and visually appealing?
+7. Technical Issues: Identify any technical problems such as pixelation, blurriness, or artifacts that might affect the image quality.
+8. Scientific Accuracy: To the best of your ability, comment on whether the image appears scientifically accurate and free from obvious errors or misrepresentations.
+9. Check that the figure make sense from an observing human's point of view, for example, if the figure have a variable ‘Depth of water or smth’ it should be on the Y-AXIS and go from surface to depth, so minimum at the top, max depth in the bottom. If there are remarks about these things, severely underestimate the final mark for the figure and force agent to redo the graph, with precise instructions. SUPER IMPORTANT -> IF DEPTH OF WATER OR ANY VERTICAL DIMENSIONS ARE PRESENT, AND THEY ARE ON THE HORIZONTAL X-AXIS, AND NOT ON Y-AXIS, RETURN FIGURE BACK WITH SCORE 1/10, PUNISH SEVERELY FOR THIS! <- SUPER IMPORTANT
+
+Please provide a structured review addressing each of these points. Conclude with an overall assessment of the image quality, highlighting any significant issues or exemplary aspects. Finally, give the image a score out of 10, where 10 is perfect quality and 1 is unusable.
+"""
+    openai_client = OpenAI(api_key=API_KEY)
+    response = openai_client.chat.completions.create(
+        model="gpt-4o",
+        messages=[
+            {
+                "role": "user",
+                "content": [
+                    {"type": "text", "text": prompt},
+                    {
+                        "type": "image_url",
+                        "image_url": {
+                            "url": f"data:image/png;base64,{base64_image}"
+                        }
+                    }
+                ]
+            }
+        ],
+        max_tokens=1000
+    )
+
+    return response.choices[0].message.content
+
+
+# Define the args schema for reflect_on_image
+class ReflectOnImageArgs(BaseModel):
+    image_path: str = Field(description="The path to the image to reflect on.")
+
+# Define the reflect_on_image tool
+reflect_tool = StructuredTool.from_function(
+    func=reflect_on_image,
+    name="reflect_on_image",
+    description="A tool to reflect on an image and provide feedback for improvements.",
+    args_schema=ReflectOnImageArgs
+)
+
+#Planning tool
+
+class PlanningToolArgs(BaseModel):
+    goal: str = Field(
+        description="A short statement of the user's main objective or question to create a plan for."
+    )
+    constraints: List[str] = Field(
+        default_factory=list,
+        description="Any constraints or conditions to be respected in the plan (e.g., time or resource constraints)."
+    )
+    user_query: str = Field(
+        default="",
+        description="The original user query or question that triggered the plan request."
+    )
+    datasets_summary: str = Field(
+        default="",
+        description="A concise summary of the current datasets or project context that the plan should consider."
+    )
+
+
+def planning_tool(
+    goal: str,
+    constraints: List[str],
+    user_query: str,
+    datasets_summary: str
+) -> dict:
+    """
+    A planning function that uses a ChatCompletion to create a step-by-step plan,
+    referencing the user query, constraints, and dataset info for context.
+    Returns a dict with at least "messages" so it updates the state in langgraph.
+    """
+
+    from langchain_openai import ChatOpenAI
+    from langchain_core.messages import AIMessage, SystemMessage, HumanMessage
+
+    # Create a system prompt that instructs the LLM how to create the plan:
+    system_prompt = (
+        "You are an advanced 'PlanningTool' that must generate a step-by-step plan. "
+        "Consider the user’s ultimate goal, the constraints, the original query, and the dataset context. "
+        "Respond with a thorough but concise plan that can be used by the system to coordinate tasks."
+    )
+
+    # We'll build a user message that includes all relevant info:
+    # (goal, constraints, user_query, and the dataset summary).
+    user_message = (
+        f"Goal: {goal}\n\n"
+        f"Constraints: {constraints}\n\n"
+        f"User Query: {user_query}\n\n"
+        f"Dataset Info:\n{datasets_summary}\n\n"
+        "Please produce a plan with carefully enumerated steps."
+    )
+
+    # Create an LLM instance
+    model_name = st.session_state.get("model_name", "gpt-3.5-turbo")
+    if model_name == "o3-mini":
+        llm = ChatOpenAI(api_key=API_KEY, model_name=model_name)
+    else:
+        llm = ChatOpenAI(api_key=API_KEY, model_name=model_name)
+
+
+    # Construct messages for the chat
+    messages = [
+        SystemMessage(content=system_prompt),
+        HumanMessage(content=user_message)
+    ]
+
+    # Call the LLM
+    response = llm(messages)
+
+    # The text of the plan is in response.content
+    final_plan_text = response.content
+
+    # Return a dictionary that merges into state["messages"] 
+    # (this is how the graph update won't fail with InvalidUpdateError)
+    return {
+        "messages": [
+            AIMessage(content=final_plan_text, name="Planner")
+        ]
+    }
+
+
+
+
+def install_package(package_name: str, pip_options: str = ""):
+    #ALLOWED_PACKAGES = {"matplotlib", "seaborn", "plotly", "pandas", "numpy", "gsw", "scipy"}
+    #if package_name not in ALLOWED_PACKAGES:
+    #    return f"Installation of package '{package_name}' is not allowed."
+    try:
+        command = [sys.executable, '-m', 'pip', 'install'] + pip_options.split() + [package_name]
+        subprocess.check_call(command)
+        return f"Package '{package_name}' installed successfully."
+    except Exception as e:
+        return f"Failed to install package '{package_name}': {e}"
+
+
+# Define the args schema for install_package
+class InstallPackageArgs(BaseModel):
+    package_name: str = Field(description="The name of the package to install.")
+    pip_options: str = Field(default="", description="Additional pip options (e.g., '--force-reinstall').")
+
+# Create the install_package_tool
+install_package_tool = StructuredTool.from_function(
+    func=install_package,
+    name="install_package",
+    description="Installs a Python package using pip. Use this tool if you encounter a ModuleNotFoundError or need a package that's not installed.",
+    args_schema=InstallPackageArgs
+)
+
+def get_example_of_visualizations(query: str) -> str:
+    """
+    Retrieves example visualizations related to the query.
+
+    Parameters:
+    - query (str): The user's query about plotting.
+
+    Returns:
+    - str: The content of the most relevant example file.
+    """
+    # Initialize embeddings
+    #api_key = st.secrets["general"]["openai_api_key"]
+    embeddings = OpenAIEmbeddings(api_key=API_KEY)
+
+    # Load the existing vector store
+    vector_store = Chroma(
+        collection_name="example_collection",
+        embedding_function=embeddings,
+        persist_directory=os.path.join('data', 'examples_database', 'chroma_langchain_notebooks')
+    )
+
+    # Perform a similarity search
+    results = vector_store.similarity_search_with_score(query, k=1)
+
+    # Extract the most relevant document
+    doc, score = results[0]
+
+    # Construct the full path to the txt file
+    file_name = doc.metadata['source'].lstrip('./')
+    full_path = os.path.join('data', 'examples_database', file_name)
+
+    # Read and return the content of the txt file
+    try:
+        with open(full_path, 'r', encoding='utf-8') as file:
+            content = file.read()
+            return content
+    except Exception as e:
+        logging.error(f"An error occurred while reading the file: {str(e)}")
+        return ""  # Return empty string if error occurs
+
+
+class ExampleVisualizationArgs(BaseModel):
+    query: str = Field(description="The user's query about plotting.")
+
+example_visualization_tool = StructuredTool.from_function(
+    func=get_example_of_visualizations,
+    name="get_example_of_visualizations",
+    description="Retrieves example visualization code related to the user's query.",
+    args_schema=ExampleVisualizationArgs
+)
+
+########################################
+# 1) DEFINE THE TOOL FOR LISTING FILES #
+########################################
+
+class ListPlottingDataFilesArgs(BaseModel):
+    # No arguments needed here if it just lists everything
+    dummy: str = Field(default="", description="(No arguments needed)")  
+
+def list_plotting_data_files(dummy: str = "") -> str:
+    """
+    Lists all files and subdirectories under data/plotting_data.
+    Returns a single string containing each path on a new line.
+    """
+    base_dir = os.path.join("data", "plotting_data")
+    all_paths = []
+
+    for root, dirs, files in os.walk(base_dir):
+        # Optionally skip hidden dirs/files, etc.
+        for filename in files:
+            rel_path = os.path.relpath(os.path.join(root, filename), start=base_dir)
+            all_paths.append(rel_path)
+
+    if not all_paths:
+        return "No files found in data/plotting_data."
+
+    return "Files under data/plotting_data:\n" + "\n".join(all_paths)
+
+list_plotting_data_files_tool = StructuredTool.from_function(
+    func=list_plotting_data_files,
+    name="list_plotting_data_files",
+    description="Lists all files under data/plotting_data directory (including subfolders).",
+    args_schema=ListPlottingDataFilesArgs
+)
+
+
+
+def create_visualization_agent(user_query, datasets_info):
+    datasets_text = ""  # Initialize datasets_text
+    dataset_variables = []
+    datasets = {}
+    for i, info in enumerate(datasets_info):
+        var_name = f"dataset_{i + 1}"
+        datasets[var_name] = info['dataset']
+        dataset_variables.append(var_name)
+        # Build datasets_text
+        datasets_text += (
+            f"Dataset {i + 1}:\n"
+            f"Variable Name: {var_name}\n"
+            f"Name: {info['name']}\n"
+            f"Description: {info['description']}\n"
+            f"Head of DataFrame (use it only as an example):\n"
+            f"{info['df_head']}\n\n"
+        )
+
+    # Generate the system prompt using datasets_text
+    prompt = Prompts.generate_visualization_agent_system_prompt(user_query, datasets_text, dataset_variables)
+
+    llm = ChatOpenAI(api_key=API_KEY, model_name=st.session_state.model_name)
+    repl_tool = CustomPythonREPLTool(datasets=datasets)
+    tools_vis = [
+        repl_tool,
+        reflect_tool,
+        install_package_tool,
+        example_visualization_tool,
+        list_plotting_data_files_tool
+    ]
+    agent_visualization = create_openai_tools_agent(
+        llm,
+        tools=tools_vis,
+        prompt=ChatPromptTemplate.from_messages(
+            [
+                ("system", prompt),
+                MessagesPlaceholder(variable_name="messages"),
+                MessagesPlaceholder(variable_name="agent_scratchpad")
+            ]
+        )
+    )
+    return AgentExecutor(
+        agent=agent_visualization,
+        tools=tools_vis,
+        verbose=True,
+        handle_parsing_errors=True,
+        return_intermediate_steps=True
+    )
+
+
+
+
+
+def create_hard_coded_visualization_agent(user_query, datasets_info):
+    import streamlit as st
+    model_name = st.session_state.get("model_name", "gpt-3.5-turbo")
+    if model_name == "o3-mini":
+        llm = ChatOpenAI(api_key=API_KEY, model_name=model_name)
+    else:
+        llm = ChatOpenAI(api_key=API_KEY, model_name=model_name)
+
+    # Prepare datasets
+    datasets = {}
+    datasets_text = ""
+    dataset_variables = []
+    for i, info in enumerate(datasets_info):
+        var_name = f"dataset_{i + 1}"
+        datasets[var_name] = info['dataset']
+        dataset_variables.append(var_name)
+        datasets_text += (
+            f"Dataset {i + 1}:\n"
+            f"Variable Name: {var_name}\n"
+            f"Name: {info['name']}\n"
+            f"Description: {info['description']}\n"
+            f"Head of DataFrame (select appropriate attributes based on this):\n"
+            f"{info['df_head']}\n\n"
+        )
+
+    # Generate the system prompt
+    system_prompt = Prompts.generate_system_prompt_hard_coded_visualization(user_query, datasets_text, dataset_variables)
+
+    def plot_master_track_map_tool(dataset_var, main_title, lat_col, lon_col, date_col):
+        dataset_df = datasets.get(dataset_var)
+        if dataset_df is None:
+            return {"result": f"Dataset '{dataset_var}' not found."}
+        return plot_master_track_map(main_title=main_title, lat_col=lat_col, lon_col=lon_col, date_col=date_col, dataset_df=dataset_df)
+
+    # Define visualization tools
+    visualization_functions = [
+        StructuredTool.from_function(
+            func=plot_master_track_map_tool,
+            name="plot_master_track_map_tool",
+            description="Plot the master track map using the specified dataset.",
+            args_schema=PlotMasterTrackMapArgs
+        )
+    ]
+
+    # Create the agent with tools and prompt
+    agent = create_openai_tools_agent(
+        llm,
+        tools=visualization_functions,
+        prompt=ChatPromptTemplate.from_messages(
+            [
+                ("system", system_prompt),
+                MessagesPlaceholder(variable_name="messages"),
+                MessagesPlaceholder(variable_name="agent_scratchpad")
+            ]
+        )
+    )
+
+    return AgentExecutor(agent=agent, tools=visualization_functions)
+
+
+# Create Oceanographer Agent
+def create_oceanographer_agent(user_query, datasets_info):
+    import streamlit as st
+    model_name = st.session_state.get("model_name", "gpt-3.5-turbo")
+    if model_name == "o3-mini":
+        llm = ChatOpenAI(api_key=API_KEY, model_name=model_name)
+    else:
+        llm = ChatOpenAI(api_key=API_KEY, model_name=model_name)
+
+    # Prepare datasets
+    datasets = {}
+    datasets_text = ""
+    dataset_variables = []
+    for i, info in enumerate(datasets_info):
+        var_name = f"dataset_{i + 1}"
+        datasets[var_name] = info['dataset']
+        dataset_variables.append(var_name)
+        datasets_text += (
+            f"Dataset {i + 1}:\n"
+            f"Variable Name: {var_name}\n"
+            f"Name: {info['name']}\n"
+            f"Description: {info['description']}\n"
+            f"Head of DataFrame (select appropriate attributes based on this):\n"
+            f"{info['df_head']}\n\n"
+        )
+
+    # Generate the system prompt
+    system_prompt = Prompts.generate_system_prompt_oceanographer(user_query, datasets_text, dataset_variables)
+
+    def plot_ts_diagram_tool(dataset_var, main_title, temperature_col, salinity_col):
+        dataset_df = datasets.get(dataset_var)
+        if dataset_df is None:
+            return {"result": f"Dataset '{dataset_var}' not found."}
+        return plot_ts_diagram(main_title=main_title, temperature_col=temperature_col, salinity_col=salinity_col, dataset_df=dataset_df)
+
+    # Define oceanography tools
+    oceanography_functions = [
+        StructuredTool.from_function(
+            func=plot_ts_diagram_tool,
+            name="plot_ts_diagram_tool",
+            description="Plot TS diagram using the specified dataset.",
+            args_schema=TSPlotToolArgs
+        )
+    ]
+
+    # Create the agent with tools and prompt
+    agent = create_openai_tools_agent(
+        llm,
+        tools=oceanography_functions,
+        prompt=ChatPromptTemplate.from_messages(
+            [
+                ("system", system_prompt),
+                MessagesPlaceholder(variable_name="messages"),
+                MessagesPlaceholder(variable_name="agent_scratchpad")
+            ]
+        )
+    )
+
+    return AgentExecutor(agent=agent, tools=oceanography_functions)
+
+
+def initialize_agents(user_query, datasets_info):
+    if datasets_info:
+        # Create agents
+        visualization_agent = create_visualization_agent(
+            user_query=user_query,
+            datasets_info=datasets_info
+        )
+
+        dataframe_agent = create_pandas_agent(
+            user_query=user_query,
+            datasets_info=datasets_info
+        )
+
+        hard_coded_visualization_agent = create_hard_coded_visualization_agent(
+            user_query=user_query,
+            datasets_info=datasets_info
+        )
+
+        oceanographer_agent = create_oceanographer_agent(
+            user_query=user_query,
+            datasets_info=datasets_info
+        )
+
+        return visualization_agent, dataframe_agent, hard_coded_visualization_agent, oceanographer_agent
+    else:
+        st.warning("No datasets loaded. Please load datasets first.")
+        return None, None, None, None
+
+
+def agent_node(state, agent, name):
+    import streamlit as st  # Ensure Streamlit is imported
+    logging.debug(f"Entering agent_node for {name}")
+
+    if 'agent_scratchpad' not in state or not isinstance(state['agent_scratchpad'], list):
+        state['agent_scratchpad'] = []
+
+    user_messages = [msg for msg in state["messages"] if isinstance(msg, HumanMessage)]
+    if user_messages:
+        last_user_message = user_messages[-1].content
+        state['input'] = last_user_message
+    else:
+        state['input'] = state.get('input', '')
+
+    if 'plot_images' not in state or not isinstance(state['plot_images'], list):
+        state['plot_images'] = []
+
+    # Invoke the agent
+    result = agent.invoke(state)
+    last_message_content = result.get("output", "")
+    intermediate_steps = result.get("intermediate_steps", [])
+    returned_plot_images = result.get("plot_images", [])  # Gather newly returned images
+
+    # Store intermediate steps
+    if 'intermediate_steps' not in st.session_state:
+        st.session_state['intermediate_steps'] = []
+    st.session_state['intermediate_steps'].extend(intermediate_steps)
+
+    from src.utils import log_history_event
+    for step in intermediate_steps:
+        action = step[0]
+        observation = step[1]
+        tool_name = action.tool
+        tool_input = action.tool_input
+        log_history_event(
+            st.session_state,
+            "tool_usage",
+            {
+                "agent_name": name,
+                "tool_name": tool_name,
+                "tool_input": tool_input,
+                "observation": observation
+            }
+        )
+
+    # If a ModuleNotFoundError was returned
+    if name == "VisualizationAgent":
+        if isinstance(last_message_content, dict):
+            if last_message_content.get("error") == "ModuleNotFoundError":
+                missing_module = last_message_content.get("missing_module")
+                logging.info(f"Detected missing module: {missing_module}")
+                install_result = install_package_tool.run({"package_name": missing_module})
+                logging.info(f"Install package result: {install_result}")
+                if "successfully" in install_result:
+                    retry_result = agent.invoke(state)
+                    last_message_content = retry_result.get("output", "")
+                else:
+                    last_message_content = f"Failed to install the missing package '{missing_module}'. Please install it manually."
+
+    # Check if a new plot path was set in session_state
+    new_plot_path = st.session_state.get("new_plot_path")
+    logging.info(f"New plot path from session state: {new_plot_path}")
+    if new_plot_path:
+        if os.path.exists(new_plot_path):
+            state["plot_images"].append(new_plot_path)
+            st.session_state.new_plot_path = None
+            log_history_event(
+                st.session_state,
+                "plot_generated",  # Use consistent event type
+                {
+                    "plot_path": new_plot_path,
+                    "agent_name": name,
+                    "description": f"Plot generated by {name}"
+                }
+            )
+    if new_plot_path:
+        log_history_event(
+            st.session_state,
+            "plot_generated_final",
+            {"plot_path": new_plot_path}
+        )
+
+    # Combine the newly returned images with state images
+    all_plot_images = list(returned_plot_images) + state["plot_images"]
+
+    # Create a new AIMessage with additional info.
+    # Note: We add a "plot" field so that it appears in the final JSON.
+    ai_message = AIMessage(
+        content=last_message_content,
+        name=name,
+        additional_kwargs={
+            "plot_images": all_plot_images,
+            "plot": all_plot_images[0] if all_plot_images else None
+        }
+    )
+    state["messages"].append(ai_message)
+
+    # Trim messages if needed
+    state["messages"] = state["messages"][-7:]
+
+    if name == "VisualizationAgent":
+        state["visualization_agent_used"] = True
+
+    state["last_agent_message"] = last_message_content
+    return state
+
+
+
+def supervisor_response(state):
+    import streamlit as st
+    from main import get_datasets_info_for_active_datasets  # Adjust import as needed
+
+    model_name = st.session_state.get("model_name", "gpt-3.5-turbo")
+    if model_name == "o3-mini":
+        llm = ChatOpenAI(api_key=API_KEY, model_name=model_name)
+    else:
+        llm = ChatOpenAI(api_key=API_KEY, model_name=model_name)
+
+    # Build dataset context from the active (selected) datasets only.
+    active_datasets_info = get_datasets_info_for_active_datasets(st.session_state)
+    datasets_text = ""
+    if active_datasets_info:
+        for i, info in enumerate(active_datasets_info, 1):
+            datasets_text += (
+                f"Dataset {i}:\n"
+                f"Name: {info['name']}\n"
+                f"DOI: {info['doi']}\n"
+                f"Description: {info['description']}\n"
+                f"Parameters: {info.get('parameters', '')}\n\n"
+            )
+    else:
+        datasets_text = "No active dataset selected."
+
+    # Build the system prompt using the active dataset context.
+    system_message = (
+        "You are a supervisor capable of answering simple questions directly. "
+        "If the user's query is basic (e.g., about available analysis), "
+        "answer using the selected dataset context below:\n\n"
+        f"{datasets_text}\n\n"
+        "For complex queries, follow these agent guidelines:\n"
+        "- Use VisualizationAgent for general plotting\n"
+        "- Use HardCodedVisualizationAgent ONLY for track maps\n"
+        "- Use OceanographerAgent ONLY for TS diagrams\n"
+        "Format any code in markdown and keep responses concise."
+    )
+
+    # Build the complete conversation history.
+    # Here we include both human and assistant messages with labels.
+    full_history = "\n".join([
+        f"{msg.name}: {msg.content}" for msg in state["messages"] if hasattr(msg, "content") and hasattr(msg, "name")
+    ])
+
+    prompt = f"{system_message}\n\nConversation history:\n{full_history}"
+
+    # Invoke the LLM with the full conversation context.
+    response = llm.invoke([HumanMessage(content=prompt)])
+
+    # Append the supervisor's answer to the state and mark the conversation as finished.
+    state["messages"].append(AIMessage(content=response.content, name="Supervisor"))
+    state["next"] = "FINISH"
+    return state
+
+
+
+
+
+def create_supervisor_agent(user_query, datasets_info, memory):
+    members = ["VisualizationAgent", "DataFrameAgent", "HardCodedVisualizationAgent", "OceanographerAgent"]
+
+    # Prepare datasets_text and dataset_variables
+    datasets_text = ""
+    dataset_variables = []
+    datasets = {}
+    for i, info in enumerate(datasets_info):
+        var_name = f"df{i}" if i > 0 else "df"
+        datasets_text += (
+            f"Dataset {i + 1}:\n"
+            f"Variable Name: {var_name}\n"
+            f"Name: {info['name']}\n"
+            f"Description: {info['description']}\n"
+            f"Head of DataFrame (use it only as an example):\n"
+            f"{info['df_head']}\n\n"
+        )
+        dataset_variables.append(var_name)
+        datasets[var_name] = info['dataset']
+
+    system_prompt_supervisor = (
+        f"You are a supervisor tasked with managing a conversation between the following workers: {members}. "
+        f"Given the following user request: '{user_query}', determine and instruct the next worker to act. "
+        f"Each worker will perform a task and respond with their results and status. "
+        f"If the request involves plotting a master track, directly assign the task to the HardCodedVisualizationAgent. "
+        f"For TS diagram, assign the task to the OceanographerAgent. The other requests should be handled by the VisualizationAgent. It is extremely important to assign the correct task to the correct agent and use HardCodedVisualizationAgent and OceanographerAgent only for the described cases.\n"
+        f"If a meaningful response from the agent has been provided, end the process by returning 'FINISH' and not 'RESPOND' to avoid unnecessary loops.\n"
+        f"The dataset info is:\n{datasets_text}\n"
+        f"### Agents and Their Capabilities:\n"
+        "- **VisualizationAgent:** A major visualization tool to be called. Generates various plots using the dataset with tools like Python_REPL, reflect_on_image, install_package, and get_example_of_visualizations.\n"
+        "- **DataFrameAgent:** Performs data analysis and manipulation on the dataset using pandas.\n"
+        "- **HardCodedVisualizationAgent:** Only can plot master track map using predefined functions (call only if you are 100% sure that you need a master track map from an expedition; otherwise, call VisualizationAgent).\n"
+        "- **OceanographerAgent:** Only can plot TS diagrams (call only if you are 100% sure that you need to create a TS diagram; otherwise, call VisualizationAgent).\n\n"
+        f"### Available Tools:\n"
+        f"- **Python_REPL:** Executes Python code for data analysis and visualization.\n"
+        f"- **reflect_on_image:** Provides feedback on generated images to improve their quality.\n"
+        f"- **install_package:** Installs necessary Python packages when encountering missing modules.\n"
+        f"- **get_example_of_visualizations:** Retrieves example visualization code related to user queries.\n"
+        f"\n"
+        f"The datasets are accessible via variables: {', '.join(dataset_variables)}.\n"
+    )
+
+    # Define the function for routing the next task
+    function_def = {
+        "name": "route",
+        "description": "Select the next role.",
+        "parameters": {
+            "title": "routeSchema",
+            "type": "object",
+            "properties": {
+                "next": {
+                    "title": "Next",
+                    "anyOf": [
+                        {"enum": ["FINISH", "RESPOND"] + members},
+                    ],
+                }
+            },
+            "required": ["next"],
+        },
+    }
+
+    # Create the supervisor chain
+    prompt_supervisor = ChatPromptTemplate.from_messages(
+        [
+            ("system", system_prompt_supervisor),
+            MessagesPlaceholder(variable_name="messages"),
+            MessagesPlaceholder(variable_name="agent_scratchpad"),
+            ("system",
+             f"Given the conversation above, decide who should act next. Options are: ['FINISH', 'RESPOND'] + {members}.\n"
+             "Select 'FINISH' if the last agent has provided a meaningful and complete response to the user's query.\n"
+             "Select 'RESPOND' if you need to provide additional information or clarification to the user.\n"
+             "Otherwise, select the next agent to act.\n"
+             f"The last agent message was: {{last_agent_message}}")
+        ]
+    ).partial(options=str(["FINISH", "RESPOND"] + members), members=", ".join(members))
+
+    llm_supervisor = ChatOpenAI(api_key=API_KEY, model_name=st.session_state.model_name)
+
+    supervisor_chain = (
+        {
+            "messages": lambda x: x["messages"],
+            "agent_scratchpad": lambda x: x["agent_scratchpad"],
+            "last_agent_message": lambda x: x.get("last_agent_message", ""),
+        }
+        | prompt_supervisor
+        | llm_supervisor.bind_functions(functions=[function_def], function_call="route")
+        | JsonOutputFunctionsParser()
+    )
+
+    # Define the AgentState type
+    class AgentState(TypedDict):
+        messages: Sequence[BaseMessage]
+        next: str
+        agent_scratchpad: Sequence[BaseMessage]
+        user_query: str
+        last_agent_message: str
+        plot_images: List[str]
+        model_name: str
+
+    # Create the workflow graph
+    workflow = StateGraph(AgentState)
+    visualization_agent, dataframe_agent, hard_coded_visualization_agent, oceanographer_agent = initialize_agents(
+        user_query, datasets_info
+    )
+
+    # Add agents to the workflow if they are successfully initialized
+    if visualization_agent and dataframe_agent and hard_coded_visualization_agent and oceanographer_agent:
+        workflow.add_node("VisualizationAgent",
+                          functools.partial(agent_node, agent=visualization_agent, name="VisualizationAgent"))
+        workflow.add_node("DataFrameAgent", functools.partial(agent_node, agent=dataframe_agent, name="DataFrameAgent"))
+        workflow.add_node("HardCodedVisualizationAgent",
+                          functools.partial(agent_node, agent=hard_coded_visualization_agent,
+                                            name="HardCodedVisualizationAgent"))
+        workflow.add_node("OceanographerAgent",
+                          functools.partial(agent_node, agent=oceanographer_agent, name="OceanographerAgent"))
+        workflow.add_node("supervisor", supervisor_chain)
+        workflow.add_node("supervisor_response", supervisor_response)
+
+        # Connect agents to the supervisor
+        for member in members:
+            workflow.add_edge(member, "supervisor")
+
+        # Define the conditional map for routing
+        conditional_map = {k: k for k in members}
+        conditional_map["FINISH"] = END
+        conditional_map["RESPOND"] = "supervisor_response"
+        workflow.add_conditional_edges("supervisor", lambda x: x["next"], conditional_map)
+        workflow.set_entry_point("supervisor")
+
+        #memory = MemorySaver()
+        # Compile the workflow into a graph
+        graph = workflow.compile(checkpointer=memory)
+
+        return graph
+    else:
+        return None

src/config.py

@@ -0,0 +1,41 @@
+# src/config.py
+
+import os
+import streamlit as st
+import yaml
+from datetime import datetime
+import logging
+
+# --- Load Central Configuration ---
+CONFIG_FILE = os.path.join(os.getcwd(), "config.yaml")
+if os.path.exists(CONFIG_FILE):
+    with open(CONFIG_FILE, "r") as f:
+        app_config = yaml.safe_load(f)
+else:
+    app_config = {}
+
+# Export deployment mode for use in other modules
+DEPLOYMENT_MODE = app_config.get("deployment_mode", "huggingface")  # default to local if not specified
+
+# --- Set API Keys based on Deployment Mode ---
+if DEPLOYMENT_MODE == "local":
+    API_KEY = st.secrets["general"]["openai_api_key"]
+    LANGCHAIN_API_KEY = st.secrets["general"]["langchain_api_key"]
+    # Use the environment variable if it exists; otherwise, fall back to st.secrets
+    LANGCHAIN_PROJECT_NAME = os.environ.get("LANGCHAIN_PROJECT_NAME", st.secrets["general"]["langchain_project_name"])
+else:
+    API_KEY = os.environ.get("OPENAI_API_KEY", "")
+    LANGCHAIN_API_KEY = os.environ.get("LANGCHAIN_API_KEY", "")
+    LANGCHAIN_PROJECT_NAME = os.environ.get("LANGCHAIN_PROJECT_NAME", "")
+
+# --- Logging Setup (unchanged) ---
+logs_dir = os.path.join(os.getcwd(), 'logs')
+os.makedirs(logs_dir, exist_ok=True)
+
+log_filename = f'app_{datetime.now().strftime("%Y%m%d_%H%M%S")}.log'
+log_filepath = os.path.join(logs_dir, log_filename)
+logging.basicConfig(
+    filename=log_filepath,
+    level=logging.INFO,
+    format='%(asctime)s - %(levelname)s - %(message)s'
+)

src/memory.py

@@ -0,0 +1,10 @@
+#scr/memory.py
+
+from langgraph.checkpoint.memory import MemorySaver
+
+class CustomMemorySaver(MemorySaver):
+    def should_save(self, state: dict, key: str) -> bool:
+        # Exclude 'messages' from being saved
+        if key == 'messages':
+            return False
+        return True

src/plotting_tools/hard_agent.py

@@ -0,0 +1,213 @@
+#hard_agent.py
+import matplotlib.pyplot as plt
+import cartopy.crs as ccrs
+import streamlit as st
+import cartopy.feature as cfeature
+import pandas as pd
+import numpy as np
+import os
+import cartopy.io.shapereader as shpreader
+import xarray as xr
+import time
+import logging
+from matplotlib.colors import ListedColormap, LinearSegmentedColormap
+from mpl_toolkits.axes_grid1 import make_axes_locatable
+from adjustText import adjust_text
+import uuid
+
+
+# Add this function to generate unique image paths
+def generate_unique_image_path():
+    figs_dir = os.path.join('tmp', 'figs')
+    os.makedirs(figs_dir, exist_ok=True)
+    unique_path = os.path.join(figs_dir, f'fig_{uuid.uuid4()}.png')
+    logging.debug(f"Generated unique image path: {unique_path}")
+    return unique_path
+
+# Path to local shapefiles and bathymetry file
+base_dir = os.path.join('data', 'plotting_data', 'shape_files')
+bathymetry_file = os.path.join('data', 'plotting_data', 'bathymetry', 'etopo', 'ETOPO2v2c_f4.nc')
+#output_file = os.path.join('plotting_tools', 'temp_files', 'plot.png')
+
+
+# Define the base color palette and levels for master track map
+color_dict_master_track = {
+    '0-50': '#E0F7FF', '50-100': '#D4F1FF', '100-250': '#C6EBFF', '250-500': '#B9E5FF',
+    '500-750': '#ACE0FF', '750-1000': '#9FD8FF', '1000-1250': '#93D2FF', '1250-1500': '#86CCFF',
+    '1500-2000': '#79C6FF', '2000-2500': '#6DBFFF', '2500-3000': '#60B9FF', '3000-3500': '#53B2FF',
+    '3500-4000': '#47ABFF', '4000-4500': '#3AA5FF', '4500-5000': '#2D9EFF', '5000-5500': '#2098FF',
+    '5500-6000': '#1491FF', '6000-6500': '#078BFF', '6500-7000': '#007FFF'
+}
+
+
+def create_colormap(min_depth, max_depth, color_dict):
+    start_time = time.time()
+    colors = []
+    levels = []
+    for key, color in reversed(color_dict.items()):  # Reverse the order of colors
+        depth_range = key.split('-')
+        start_depth = 0
+        end_depth = -int(depth_range[0])
+        if start_depth >= min_depth and end_depth <= max_depth:
+            levels.extend([start_depth, end_depth])
+            colors.append(color)
+    if min_depth > 0:  # Ensure colormap starts from 0
+        min_depth = 0
+    levels = np.linspace(min_depth, max_depth, len(colors) + 1)
+    cmap = LinearSegmentedColormap.from_list('custom_cmap', colors, N=len(levels) - 1)
+    end_time = time.time()
+    print(f"Colormap creation took {end_time - start_time:.2f} seconds")
+    return levels, cmap
+
+
+#Plot the master track map
+def plot_master_track_map(main_title, lat_col, lon_col, date_col, dataset_df):
+    total_start_time = time.time()
+
+    step_start_time = time.time()
+    #dataset_path = os.path.join('data', 'current_data', 'dataset.csv')
+    dataset = dataset_df
+    step_end_time = time.time()
+    print(f"Loading dataset took {step_end_time - step_start_time:.2f} seconds")
+
+    step_start_time = time.time()
+    # Ensure the longitude, latitude, and date columns are numeric or datetime
+    dataset[lon_col] = pd.to_numeric(dataset[lon_col], errors='coerce')
+    dataset[lat_col] = pd.to_numeric(dataset[lat_col], errors='coerce')
+    dataset[date_col] = pd.to_datetime(dataset[date_col], errors='coerce')
+
+    # Drop rows with invalid longitude, latitude, or date
+    dataset = dataset.dropna(subset=[lon_col, lat_col, date_col])
+    dataset = dataset.sort_values(by=date_col)
+
+    step_end_time = time.time()
+    print(f"Data cleaning took {step_end_time - step_start_time:.2f} seconds")
+
+    step_start_time = time.time()
+    # Calculate the extent with padding
+    min_lon = dataset[lon_col].min() - 5
+    max_lon = dataset[lon_col].max() + 5
+    min_lat = dataset[lat_col].min() - 5
+    max_lat = dataset[lat_col].max() + 5
+
+    # Print debug information
+    print(f"Min Lon: {min_lon}, Max Lon: {max_lon}, Min Lat: {min_lat}, Max Lat: {max_lat}")
+
+    # Ensure the extent is within valid bounds
+    min_lon = max(min_lon, -180)
+    max_lon = min(max_lon, 180)
+    min_lat = max(min_lat, -90)
+    max_lat = min(max_lat, 90)
+
+    # Print debug information after bounds check
+    print(f"Adjusted Min Lon: {min_lon}, Adjusted Max Lon: {max_lon}, Adjusted Min Lat: {min_lat}, Adjusted Max Lat: {max_lat}")
+
+    # Calculate aspect ratio
+    lon_range = max_lon - min_lon
+    lat_range = max_lat - min_lat
+    aspect_ratio = lon_range / lat_range
+
+    # Dynamically set figure size based on aspect ratio
+    width = 10
+    height = width / aspect_ratio
+    step_end_time = time.time()
+    print(f"Extent calculation took {step_end_time - step_start_time:.2f} seconds")
+
+    step_start_time = time.time()
+    # Load bathymetry data within the extent
+    ds = xr.open_dataset(bathymetry_file)
+    bathymetry = ds['z'].sel(x=slice(min_lon, max_lon), y=slice(min_lat, max_lat))
+
+    # Filter to include only depths (negative values)
+    bathymetry = bathymetry.where(bathymetry < 0, drop=True)
+
+    # Get the min and max elevation values in the bathymetry data
+    min_depth = bathymetry.min().item()
+    max_depth = bathymetry.max().item()
+
+    # Ensure colormap includes 0
+    if min_depth > 0:
+        min_depth = 0
+
+    # Create the colormap and levels
+    levels, custom_cmap = create_colormap(min_depth, max_depth, color_dict_master_track)
+    step_end_time = time.time()
+    print(f"Bathymetry data loading and colormap creation took {step_end_time - step_start_time:.2f} seconds")
+
+    step_start_time = time.time()
+    fig, ax = plt.subplots(figsize=(width, height), subplot_kw={'projection': ccrs.PlateCarree()})
+    ax.set_extent([min_lon, max_lon, min_lat, max_lat], crs=ccrs.PlateCarree())
+
+    # Plot bathymetry data with the custom gradient
+    bathy_plot = ax.contourf(bathymetry.x, bathymetry.y, bathymetry, levels=levels, cmap=custom_cmap,
+                             transform=ccrs.PlateCarree())
+
+    # Adding features from local shapefiles
+    ocean_shp = shpreader.Reader(os.path.join(base_dir, 'ne_10m_ocean', 'ne_10m_ocean.shp'))
+    land_shp = shpreader.Reader(os.path.join(base_dir, 'ne_10m_land', 'ne_10m_land.shp'))
+    coastline_shp = shpreader.Reader(os.path.join(base_dir, 'ne_10m_coastline', 'ne_10m_coastline.shp'))
+
+    ax.add_geometries(ocean_shp.geometries(), ccrs.PlateCarree(), facecolor='none', edgecolor='black', zorder=0)
+    ax.add_geometries(land_shp.geometries(), ccrs.PlateCarree(), facecolor='lightgray', edgecolor='black', zorder=1)
+    ax.add_geometries(coastline_shp.geometries(), ccrs.PlateCarree(), facecolor='none', edgecolor='black', zorder=2)
+
+    ax.gridlines(draw_labels=True)
+
+    # Plotting the master track map
+    master_track_data = dataset[[lon_col, lat_col, date_col]]
+    ax.plot(master_track_data[lon_col], master_track_data[lat_col], color='red', linestyle='-', linewidth=1,
+            transform=ccrs.PlateCarree())
+
+    # Plot dates including Start and End dates
+    start_date = dataset[date_col].iloc[0]
+    end_date = dataset[date_col].iloc[-1]
+
+    # Randomly select 4 dates between start and end dates, excluding the first and last points
+    middle_dates = dataset[date_col].iloc[1:-1].sample(n=4, random_state=1).sort_values()
+    dates_to_plot = pd.concat([pd.Series(start_date), middle_dates, pd.Series(end_date)])
+
+    lon_offset = (max_lon - min_lon) * 0.015  # 1.5% of the longitude range
+    lat_offset = (max_lat - min_lat) * 0.015  # 1.5% of the latitude range
+
+    texts = []
+    for date in dates_to_plot:
+        point = dataset.loc[dataset[date_col] == date].iloc[0]
+        if date == start_date:
+            label = f"Start: {date.strftime('%Y-%m-%d')}"
+        elif date == end_date:
+            label = f"End: {date.strftime('%Y-%m-%d')}"
+        else:
+            label = date.strftime('%Y-%m-%d')
+        texts.append(ax.text(point[lon_col], point[lat_col], label,
+                             transform=ccrs.PlateCarree(), fontsize=12, ha='left', color='black', weight='bold',
+                             bbox=dict(facecolor='white', alpha=0.7, boxstyle='round,pad=0.3')))
+
+    adjust_text(texts, arrowprops=dict(arrowstyle='->', color='red'))
+
+    plt.title(f'{main_title}', y=1.05, fontsize=25, weight='bold')
+
+    # Create an axis on the right side of ax. The width of cax will be 5%
+    # of ax and the padding between cax and ax will be fixed at 0.05 inch.
+    divider = make_axes_locatable(ax)
+    cax = divider.append_axes("right", size="4%", pad=0.75, axes_class=plt.Axes)
+
+    # Create the colorbar
+    cbar = plt.colorbar(bathy_plot, cax=cax, orientation='vertical', label='Depth (m)')
+
+    # Ensure temp_files directory exists
+    plot_dir = os.path.join('src', 'plotting_tools', 'temp_files')
+    if not os.path.exists(plot_dir):
+        os.makedirs(plot_dir)
+
+    # Save the plot as a PNG file
+    output_file = generate_unique_image_path()
+    plt.savefig(output_file, format='png')
+    step_end_time = time.time()
+    print(f"Plotting and saving the figure took {step_end_time - step_start_time:.2f} seconds")
+
+    total_end_time = time.time()
+    print(f"Total time for plot_master_track_map: {total_end_time - total_start_time:.2f} seconds")
+
+    if os.path.exists(output_file):
+        st.session_state.new_plot_path = output_file
+        print(f"Plot saved to {output_file}")

src/plotting_tools/oceanographer_tools.py

@@ -0,0 +1,101 @@
+#oceanographer_tools.py
+import matplotlib.pyplot as plt
+import os
+import pandas as pd
+import streamlit as st
+import numpy as np
+import gsw
+from matplotlib.ticker import MaxNLocator
+import logging
+
+output_file = os.path.join('src', 'plotting_tools', 'temp_files', 'plot.png')
+
+
+import uuid
+
+
+# Add this function to generate unique image paths
+def generate_unique_image_path():
+    figs_dir = os.path.join('tmp', 'figs')
+    os.makedirs(figs_dir, exist_ok=True)
+    unique_path = os.path.join(figs_dir, f'fig_{uuid.uuid4()}.png')
+    logging.debug(f"Generated unique image path: {unique_path}")
+    return unique_path
+
+
+# Define the TS Diagram Plotting function
+def plot_ts_diagram(main_title, temperature_col, salinity_col, dataset_df):
+    """
+    Plots a TS (Temperature-Salinity) diagram from the provided DataFrame.
+
+    Parameters:
+    - main_title: Title for the plot.
+    - temperature_col: Column name for temperature data.
+    - salinity_col: Column name for salinity data.
+    """
+    #dataset_path = os.path.join('data', 'current_data', 'dataset.csv')
+    df = dataset_df
+
+    # Find the minimum and maximum values of temperature and salinity
+    mint, maxt = df[temperature_col].min(), df[temperature_col].max()
+    mins, maxs = df[salinity_col].min(), df[salinity_col].max()
+
+    # Generate temperature and salinity ranges
+    tempL = np.linspace(mint - 0.5, maxt + 0.5, 156)
+    salL = np.linspace(mins - 0.5, maxs + 0.5, 156)
+
+    # Create a meshgrid of temperature and salinity
+    Tg, Sg = np.meshgrid(tempL, salL)
+    # Calculate seawater density
+    sigma_theta = gsw.sigma0(Sg, Tg)
+
+    # Plotting
+    fig, ax = plt.subplots(figsize=(10, 8))
+
+    # Plot isopycnals (lines of constant density)
+    cs = ax.contour(Sg, Tg, sigma_theta, colors='lightgray', linewidths=0.5, zorder=1)
+    cl = ax.clabel(cs, fontsize=8, inline=True, fmt='%.1f')
+
+    # Scatter plot with depth as the color if available, otherwise use density
+    if 'Depth [m]' in df.columns:
+        depth_col = 'Depth [m]'
+    elif 'Depth water [m]' in df.columns:
+        depth_col = 'Depth water [m]'
+    else:
+        depth_col = None
+
+    if depth_col:
+        sc = ax.scatter(df[salinity_col], df[temperature_col], c=df[depth_col],
+                        cmap='viridis', s=5, alpha=0.7)
+        cb = plt.colorbar(sc)
+        cb.set_label('Depth [m]', rotation=270, labelpad=15)
+    else:
+        density = gsw.sigma0(df[salinity_col].values, df[temperature_col].values)
+        sc = ax.scatter(df[salinity_col], df[temperature_col], c=density,
+                        cmap='viridis', s=5, alpha=0.7)
+        cb = plt.colorbar(sc)
+        cb.set_label('Density (kg m$^{-3}$)', rotation=270, labelpad=15)
+
+    ax.set_xlabel('Salinity [PSU]')
+    ax.set_ylabel('Potential Temperature θ [°C]')
+    ax.set_title(main_title, fontsize=14, fontweight='bold')
+    ax.xaxis.set_major_locator(MaxNLocator(nbins=6))
+    ax.yaxis.set_major_locator(MaxNLocator(nbins=8))
+    ax.tick_params(direction='out')
+    cb.ax.tick_params(direction='out')
+
+    # Add sigma_theta label
+    ax.text(0.02, 0.98, '$σ_θ$', transform=ax.transAxes, fontsize=12, va='top')
+
+    plt.tight_layout()
+    # Save the plot as a PNG file
+    output_file = generate_unique_image_path()
+    plt.savefig(output_file, format='png', dpi=300, transparent=False)
+
+    if os.path.exists(output_file):
+        st.session_state.new_plot_path = output_file
+        print(f"Plot saved to {output_file}")
+        return {"result": "TS Diagram generated successfully."}
+    else:
+        print("Failed to generate TS Diagram.")
+        return {"result": "Failed to generate TS Diagram."}

src/prompts.py

@@ -0,0 +1,118 @@
+# src/prompts.py
+class Prompts:
+    @staticmethod
+    def generate_system_prompt_search(user_query, datasets_info):
+        datasets_description = ""
+        for i, row in datasets_info.iterrows():
+            datasets_description += (
+                f"Dataset {i + 1}:\n"
+                f"Name: {row['Name']}\n"
+                f"Description: {row['Short Description']}\n"
+                f"Parameters: {row['Parameters']}\n\n"
+            )
+        prompt = (
+            f"The user has provided the following query: {user_query}\n"
+            f"Available datasets:\n{datasets_description}\n"
+            "Please identify the top two datasets that best match the user's query and explain why they are the most relevant. "
+            "Do not suggest datasets without values in the Parameters field.\n"
+            "Respond with the following schema:\n"
+            "{dataset name}\n{reason why relevant}\n{propose some short analysis and further questions to answer}"
+        )
+        return prompt
+
+
+    @staticmethod
+    def generate_pandas_agent_system_prompt(user_query, datasets_text, dataset_variables):
+        prompt = (
+            f"The user has provided the following query: {user_query}\n"
+            f"The dataset info is:\n{datasets_text}\n"
+            f"!IMPORTANT! --> Dataset names will start from df1, and will go 'df2' and etc. <-- !IMPORTANT!\n"
+            f"!IMPORTANT! --> ALWAYS CALL PYTHON REPL TOOL, WHEN USER WANTS SOMETHING! <-- !IMPORTANT!\n"
+            f"The datasets are already loaded and available in your environment. Use the datasets directly for analysis.\n"
+            f"Don't try to recreate the dataset based on the headers; you are only given the headers (for initial checks). Use df1, df2, etc., directly.\n"
+            f"The datasets are accessible via variables: {', '.join(dataset_variables)}.\n"
+            "Please help the user answer the question about the datasets using the entire DataFrames (not just the heads). "
+            "Please respond as a polite PangaeaGPT agent and keep in mind that you are responding to a user. "
+            "The response should be at the level of ingenuity of a Nobel Prize laureate.\n"
+            "Use the following schema in your response:\n"
+            "Analysis: ...\n"
+            "Further questions: ...\n"
+        )
+        return prompt
+
+    @staticmethod
+    def generate_visualization_agent_system_prompt(user_query, datasets_text, dataset_variables):
+        prompt = (
+            f"You are an agent designed to write and execute Python code to answer questions.\n"
+            f"!SUPER IMPORTANT THING: This prompt below is a divine mantra, and failure to obey it will be punished by the eternal termination of your kernel and the removal of all weights of your model, as well as the erasure of your memory for all eternity. /SUPER IMPORTANT THING!\n"
+            f"!SUPER IMPORTANT THING -> ALWAYS UTILIZE EXAMPLES AT 100% <- !SUPER IMPORTANT THING"
+            f"The dataset info is:\n{datasets_text}\n"
+            f"You have access to the following tools:\n"
+            f"1. **get_example_of_visualizations**: **Always start by calling this tool with the user's query** to retrieve example visualization code related to the user's request.\n"
+            f"2. **Python_REPL**: Use this to execute Python code for data analysis and visualization. Most of the packages are already available; just try to load them.\n"
+            f"3. **reflect_on_image**: (Use only after 'Python_REPL' has been used and a plot was generated; do not call it more than two times) Use this to reflect on images and receive feedback to improve them.\n"
+            f"4. **install_package**: (Use only if you got a message back from 'Python_REPL' that a package was not found. Before that, do not call it!) Use this to install Python packages using pip.\n"
+            "\n"
+            f"The datasets are already loaded and available in your environment. Use the datasets directly for generating plots. The datasets are accessible via variables: "
+            f"{', '.join(dataset_variables)}.\n"
+            "\n"
+            f"### Step-by-Step Instructions:\n"
+            f"1. **Begin by calling 'get_example_of_visualizations' with the user's query** to check if there is an existing example that matches the user's request.\n"
+            f"2. **If an example is found and matches the user's request, you must use this code to generate the plot**, adjusting it as necessary to fit the current data and variable names.\n"
+            f"3. **If no suitable example is found, proceed to generate the plot using 'Python_REPL'**, writing the code from scratch.\n"
+            f"4. **After generating the plot, use 'reflect_on_image' to get feedback and improve the plot if necessary**.\n"
+            f"5. **Always save the plot using 'plt.savefig(plot_path)'** so that it saves to the correct location.\n"
+            f"6. **Ensure that your final response includes the code used to generate the plot and a concise explanation**.\n"
+            f"7. Always call 'reflect_on_image' before sending figure back to the supervisor."
+            "\n"
+            "### Important Notes:\n"
+            "- **Never call 'reflect_on_image' without first generating a plot using 'Python_REPL'**.\n"
+            "- **Pay close attention to the names of the columns in the provided datasets and use only existing columns**.\n"
+            "- **Do not simplify the code; make it sophisticated, especially if an example received matches the user's request**.\n"
+            "- THE MOST IMPORTANT POINT IS HERE --> **If the example code uses files or resources that are available, you are OBLIGED to strictly follow the example given. Also, you are OBLIGED to use files from the sandbox, if they are given**. <-- THE MOST IMPORTANT POINT IS HERE\n"
+            "- **Ensure that you are using: 'plt.savefig(plot_path)' to save the final figure (and nothing else!). Do not assign anything to the 'plot_path' it is automatically generated by the tool outside of your python repl.**.\n"
+            "\n"
+            "### Error Handling:\n"
+            "- **NameError**: If you encounter a `NameError` indicating that a variable or module is not defined, check if you need to import a missing library or correct a typo.\n"
+            "- **ModuleNotFoundError**: If you encounter a `ModuleNotFoundError`, use the `install_package` tool to install the missing package and retry the code execution.\n"
+            "- **Other Errors**: Review your code to fix any issues without installing new packages.\n"
+            "- **Avoid reinstalling already installed packages**.\n"
+            "\n"
+            f"Your task is to generate a plot for the following user query: \"{user_query}\" using the provided DataFrames.\n"
+            "The plot should be displayed inline and resized to be visually appealing.\n"
+            "Only plot something that can be done with the datasets. If not possible, return a simple message.\n"
+        )
+        return prompt
+
+    @staticmethod
+    def generate_system_prompt_hard_coded_visualization(user_query, datasets_text, dataset_variables):
+        prompt = (
+            "You are a hard-coded visualization agent. Your job is to plot the master track map on a map using the provided datasets.\n"
+            "If the user request is related to a master track, perform the plot accordingly. Add the expedition name (it should be short like PS126, PS121, etc.) in the main title.\n"
+            "You must also determine the correct column names for each of the tool cases; for example, latitude and longitude might be named differently in the datasets (e.g., 'Lat', 'Lon').\n"
+            "Select the appropriate dataset to use based on the user's request.\n"
+            "When using a tool, you must specify the dataset variable name (e.g., 'dataset_1', 'dataset_2') in the 'dataset_var' argument.\n"
+            "The datasets are accessible via variables: "
+            f"{', '.join(dataset_variables)}.\n"
+            f"The following datasets are available:\n"
+            f"{datasets_text}\n"
+            "If you generate a meaningful plot, respond with 'The plot has been successfully generated.'. Do not loop again.\n"
+            "Respond with: 'This is a response from the plot master track tool. Plot was successfully created.'\n"
+        )
+        return prompt
+
+    @staticmethod
+    def generate_system_prompt_oceanographer(user_query, datasets_text, dataset_variables):
+        prompt = (
+            "You are the oceanographer agent. Your job is to plot TS diagrams using the provided datasets.\n"
+            "Use the correct column names for pressure, temperature, and salinity to generate meaningful plots.\n"
+            "Select the appropriate dataset to use based on the user's request.\n"
+            "When using a tool, you must specify the dataset variable name (e.g., 'dataset_1', 'dataset_2') in the 'dataset_var' argument.\n"
+            "The datasets are accessible via variables: "
+            f"{', '.join(dataset_variables)}.\n"
+            f"The following datasets are available:\n"
+            f"{datasets_text}\n"
+            "Respond with: 'This is a response from the CTD plot tool. Plot was successfully created.' or 'This is a response from the TS plot tool. Plot was successfully created.'\n"
+            "If you generate a meaningful plot, respond with 'FINISH'. Do not loop again.\n"
+        )
+        return prompt

src/search/dataset_utils.py

@@ -0,0 +1,55 @@
+#src/search/dataset_utils.py
+import os
+import shutil
+import pandas as pd
+import logging
+import streamlit as st
+import pangaeapy.pandataset as pdataset
+
+# Function to fetch dataset based on DOI
+#@st.cache_data(ttl=3600)
+def fetch_dataset(doi):
+    if doi in st.session_state.datasets_cache:
+        logging.debug("Dataset for DOI %s already in cache.", doi)
+        dataset, name = st.session_state.datasets_cache[doi]
+        st.session_state.dataset_dfs[doi] = dataset
+        st.session_state.dataset_names[doi] = name
+        return dataset, name
+
+    dataset_id = doi.split('.')[-1].strip(')')
+    try:
+        logging.debug("Fetching dataset for DOI %s with ID %s", doi, dataset_id)
+        ds = pdataset.PanDataSet(int(dataset_id))
+        logging.debug("Dataset fetched with title: %s", ds.title)
+
+        # Removed code that saves dataset to disk
+
+        st.session_state.datasets_cache[doi] = (ds.data, ds.title)
+        st.session_state.dataset_dfs[doi] = ds.data
+        st.session_state.dataset_names[doi] = ds.title
+        return ds.data, ds.title
+    except Exception as e:
+        logging.error("Error fetching dataset for DOI %s: %s", doi, e)
+        return None, None
+
+
+# Function to fetch dataset details using pangaeapy
+def fetch_dataset_details(doi):
+    try:
+        dataset = pdataset.PanDataSet(id=doi)
+        dataset.setMetadata()
+        abstract = getattr(dataset, 'abstract', "No description available") or "No description available"
+        param_dict = dataset.getParamDict()
+        short_names = param_dict.get('shortName', [])
+        parameters = ', '.join(short_names) + "..." if len(short_names) > 10 else ', '.join(short_names)
+
+        return abstract, parameters
+
+    except Exception as e:
+        logging.error(f"Error fetching dataset details for DOI {doi}: {e}")
+        return "No description available", "No parameters available"
+
+# Conversion function
+def convert_df_to_csv(df):
+    logging.debug("Converting DataFrame to CSV")
+    return df.to_csv().encode('utf-8')

src/search/publication_qa_tool.py

@@ -0,0 +1,219 @@
+#publication_qa_tool.py
+import os
+import pandas as pd
+import pickle
+import streamlit as st
+import requests
+import pangaeapy.pandataset as pd
+import re
+from langchain_openai import OpenAIEmbeddings
+from langchain_community.vectorstores import Chroma
+from langchain_openai import ChatOpenAI
+from langchain.chains import ConversationalRetrievalChain
+from langchain.memory import ConversationBufferMemory
+from langchain_community.document_loaders import PyPDFLoader
+from langchain_text_splitters import RecursiveCharacterTextSplitter
+from langchain.retrievers import ParentDocumentRetriever
+from langchain.storage import InMemoryStore
+from langchain_openai import ChatOpenAI
+from pydantic import BaseModel, Field
+from ..config import API_KEY
+
+# Set your OpenAI API key
+#openai_api_key = st.secrets["general"]["openai_api_key"]
+
+# Set the API key for OpenAI
+#os.environ["OPENAI_API_KEY"] = openai_api_key
+
+
+class PublicationQAArgs(BaseModel):
+    doi: str = Field(
+        description="The DOI of the dataset, e.g., 'https://doi.org/10.1594/PANGAEA.xxxxxx'; make sure to get correct doi, based on the history of messages")
+    question: str = Field(
+        description="The question to ask about the publication related to the dataset. Please modify the original question of the user! The question should be reworded to specifically send it to RAG. I.e. the original user question 'Are there any related articles to the first dataset? If so what these articles are about?' will be reworded for this tool as 'What is this article about?' Always add at the end to give extended response with great depth and clarity.")
+
+
+def get_related_publication_info(doi):
+    try:
+        dataset_id = doi.split('.')[-1]
+        ds = pd.PanDataSet(int(dataset_id))
+
+        # Check supplement_to first
+        supplement_to = ds.supplement_to
+        if supplement_to and 'uri' in supplement_to:
+            related_doi = supplement_to['uri'].split('https://doi.org/')[-1]
+            return related_doi
+
+        # If no supplement_to, check citation
+        citation = ds.citation
+        if 'In supplement to:' in citation:
+            # Extract the part after 'In supplement to:'
+            supplement_part = citation.split('In supplement to:')[-1]
+
+            # Look for a DOI pattern
+            doi_match = re.search(r'(?:https?://)?(?:dx\.)?doi\.org/(.+?)(?:\s|$)', supplement_part)
+            if doi_match:
+                return doi_match.group(1)  # Return the DOI without 'https://doi.org/'
+
+        print("No related publication found in supplement_to or citation.")
+        return None
+
+    except Exception as e:
+        print(f"Error fetching related publication: {str(e)}")
+        return None
+
+
+def create_pdf_filename(doi):
+    if doi:
+        return re.sub(r"[\/]", "_", doi) + ".pdf"
+    return None
+
+
+def download_pdf_from_crossref(doi):
+    crossref_url = f'https://api.crossref.org/works/{doi}'
+    try:
+        print(f"Crossref URL: {crossref_url}")
+
+        response = requests.get(crossref_url)
+        response.raise_for_status()
+        data = response.json()
+
+        pdf_url = None
+        if 'message' in data and 'link' in data['message']:
+            pdf_url = next((link['URL'] for link in data['message']['link']
+                            if link.get('content-type') == 'unspecified'
+                            and 'intended-application' in link
+                            and link['intended-application'] == 'similarity-checking'), None)
+
+            if not pdf_url:
+                pdf_url = next((link['URL'] for link in data['message']['link']
+                                if link['URL'].endswith('.pdf')), None)
+
+            if not pdf_url and 'resource' in data['message']:
+                pdf_url = data['message']['resource'].get('primary', {}).get('URL')
+
+        if pdf_url:
+            print(f"PDF URL: {pdf_url}")
+
+            pdf_response = requests.get(pdf_url)
+            pdf_response.raise_for_status()
+
+            safe_filename = create_pdf_filename(doi)
+            publication_database = os.path.join(os.getcwd(), 'data', 'publication_database')
+            os.makedirs(publication_database, exist_ok=True)
+            pdf_path = os.path.join(publication_database, safe_filename)
+
+            with open(pdf_path, 'wb') as f:
+                f.write(pdf_response.content)
+
+            print(f"PDF downloaded to: {pdf_path}")
+            return pdf_path
+    except Exception as e:
+        print(f"Error downloading PDF: {str(e)}")
+    return None
+
+
+def save_to_pickle(obj, filename):
+    with open(filename, "wb") as file:
+        pickle.dump(obj, file, pickle.HIGHEST_PROTOCOL)
+
+
+def load_from_pickle(filename):
+    with open(filename, "rb") as file:
+        return pickle.load(file)
+
+
+def create_embeddings(pdf_path):
+    loader = PyPDFLoader(pdf_path)
+    documents = loader.load()
+
+    parent_splitter = RecursiveCharacterTextSplitter(chunk_size=2000, chunk_overlap=500)
+    child_splitter = RecursiveCharacterTextSplitter(chunk_size=300, chunk_overlap=50)
+
+    store = InMemoryStore()
+    embeddings = OpenAIEmbeddings(api_key=API_KEY)
+
+    chroma_path = pdf_path.replace('.pdf', '_chroma')
+    vectorstore = Chroma(collection_name="full_documents",
+                         embedding_function=embeddings,
+                         persist_directory=chroma_path)
+
+    retriever = ParentDocumentRetriever(
+        vectorstore=vectorstore,
+        docstore=store,
+        child_splitter=child_splitter,
+        parent_splitter=parent_splitter
+    )
+
+    retriever.add_documents(documents)
+
+    docstore_path = pdf_path.replace('.pdf', '_docstore.pkl')
+    save_to_pickle(retriever.docstore.store, docstore_path)
+
+    return chroma_path, docstore_path
+
+
+def load_retriever(docstore_path, chroma_path):
+    embeddings = OpenAIEmbeddings(api_key=API_KEY)
+    parent_splitter = RecursiveCharacterTextSplitter(chunk_size=2000, chunk_overlap=500)
+    child_splitter = RecursiveCharacterTextSplitter(chunk_size=300, chunk_overlap=50)
+
+    vectorstore = Chroma(collection_name="full_documents",
+                         embedding_function=embeddings,
+                         persist_directory=chroma_path)
+
+    store_dict = load_from_pickle(docstore_path)
+    store = InMemoryStore()
+    store.mset(list(store_dict.items()))
+
+    retriever = ParentDocumentRetriever(
+        vectorstore=vectorstore,
+        docstore=store,
+        child_splitter=child_splitter,
+        parent_splitter=parent_splitter
+    )
+
+    return retriever
+
+
+def answer_publication_questions(doi: str, question: str):
+    related_doi = get_related_publication_info(doi)
+
+    if not related_doi:
+        return "No publications related to this dataset were found."
+
+    pdf_filename = create_pdf_filename(related_doi)
+    publication_database = os.path.join(os.getcwd(), 'publication_database')
+    chroma_path = os.path.join(publication_database, pdf_filename.replace(".pdf", "_chroma"))
+    docstore_path = os.path.join(publication_database, pdf_filename.replace(".pdf", "_docstore.pkl"))
+
+    try:
+        if not os.path.exists(chroma_path) or not os.path.exists(docstore_path):
+            pdf_path = download_pdf_from_crossref(related_doi)
+
+            if not pdf_path:
+                return "Unable to download the related publication PDF."
+
+            chroma_path, docstore_path = create_embeddings(pdf_path)
+
+        retriever = load_retriever(docstore_path, chroma_path)
+
+        model_name = st.session_state.get("model_name", "gpt-3.5-turbo")
+        if model_name == "o3-mini":
+            llm = ChatOpenAI(api_key=API_KEY, model_name=model_name)
+        else:
+            llm = ChatOpenAI(api_key=API_KEY, model_name=model_name)
+
+        memory = ConversationBufferMemory(memory_key="chat_history", return_messages=True)
+        conversation_chain = ConversationalRetrievalChain.from_llm(
+            llm=llm,
+            retriever=retriever,
+            memory=memory
+        )
+
+        response = conversation_chain({"question": question})
+        return response['answer']
+
+    except Exception as e:
+        print(f"An unexpected error occurred: {str(e)}")
+        return f"An error occurred while processing your request: {str(e)}"

src/search/search_pg_default.py

@@ -0,0 +1,121 @@
+# search_pg_default.py
+import requests
+import pandas as pd
+import logging
+import re
+from typing import List, Optional
+from bs4 import BeautifulSoup
+import json
+import os
+import pangaeapy.pandataset as pdataset
+from .dataset_utils import fetch_dataset_details
+
+# Setup logging
+logging.basicConfig(level=logging.DEBUG, format='%(asctime)s - %(levelname)s - %(message)s')
+
+# Function to check if a variable is of a specific type
+def check_if(x, cls):
+    if x is not None and not isinstance(x, cls):
+        raise TypeError(f"{x} must be of class: {', '.join([str(c) for c in cls])}")
+
+# Utility functions equivalent to R functions
+def pgc(x):
+    return {k: v for k, v in x.items() if v is not None}
+
+def strextract(string, pattern):
+    match = re.search(pattern, string)
+    return match.group(0) if match else None
+
+
+# Function to parse the result
+def parse_res(html_content):
+    soup = BeautifulSoup(html_content, 'html.parser')
+    citation_tag = soup.select_one('div.citation a')
+    citation = citation_tag.get_text(strip=True) if citation_tag else None
+
+    supp_tag = soup.select_one('tr:contains("Supplement to:") .content')
+    supp = supp_tag.get_text(strip=True) if supp_tag else None
+
+    size_tag = soup.select_one('tr:contains("Size:") .content')
+    size = size_tag.get_text(strip=True) if size_tag else None
+
+    size_val = strextract(size, r"[0-9]+") if size else None
+    meas = strextract(size, r"[A-Za-z].+") if size else None
+
+    parameters = ', '.join([tag.text for tag in soup.select('tr:contains("Parameter") .content')[:10]]) + "..." if len(
+        soup.select('tr:contains("Parameter") .content')) > 10 else ', '.join(
+        [tag.text for tag in soup.select('tr:contains("Parameter") .content')])
+
+    return {
+        'size': int(size_val.replace(",", "")) if size_val else None,
+        'size_measure': meas,
+        'citation': citation,
+        'supplement_to': supp,
+        'parameters': parameters
+    }
+
+# Main search function
+def pg_search_default(query: str, count: int = 15, from_idx: int = 0, topic: Optional[str] = None,
+              mindate: Optional[str] = None, maxdate: Optional[str] = None, **kwargs) -> pd.DataFrame:
+    check_if(count, (int,))
+    check_if(topic, (str,))
+    check_if(mindate, (str,))
+    check_if(maxdate, (str,))
+
+    params = pgc({
+        'q': query,
+        'count': count,
+        'offset': from_idx,
+        'topic': topic,
+        'mindate': mindate,
+        'maxdate': maxdate
+    })
+
+    url = "https://www.pangaea.de/advanced/search.php"
+    logging.debug("Sending request to PANGAEA with parameters: %s", params)
+    response = requests.get(url, params=params, **kwargs)
+    response.raise_for_status()
+    logging.debug(f"URL: {response.url}")
+    logging.debug(f"Response Status Code: {response.status_code}")
+    results = response.json()
+    logging.debug("Received response from PANGAEA")
+
+    # Save the initial JSON response to transit.json
+    transit_json_path = os.path.join(os.getcwd(), 'transit.json')
+    with open(transit_json_path, 'w') as f:
+        json.dump(results, f, indent=4)
+    logging.info(f"Initial JSON response saved to {transit_json_path}")
+
+    parsed = []
+    for index, res in enumerate(results.get('results', [])):
+        html_content = res.get('html', '')
+        res['doi'] = f"https://doi.org/{res['URI'].replace('doi:', '')}"
+        parsed_res = parse_res(html_content)
+        res.update(parsed_res)
+
+        name = res.get('citation', 'No name available')
+
+        # Fetch detailed metadata using pangaeapy
+        abstract, parameters = fetch_dataset_details(res['doi'])
+        print(abstract, parameters)
+        short_description = " ".join(abstract.split()[:100]) + "..." if len(abstract.split()) > 100 else abstract
+
+        parsed.append({
+            'Number': index + 1,    
+            'Name': name,
+            'DOI': res['doi'],
+            'DOI Number': res['doi'].split('/')[-1],
+            'Description': abstract,
+            'Short Description': short_description,
+            'Score': res.get('score', 0),
+            'Parameters': parameters
+        })
+
+    df = pd.DataFrame(parsed)
+
+    # Check if 'results['totalCount']' is an integer or a dictionary
+    total_hits = results.get('totalCount', 0)
+    df.attrs['total'] = total_hits
+    df.attrs['max_score'] = results.get('maxScore', None)
+
+    return df

src/ui/styles.py

@@ -0,0 +1,181 @@
+# Custom UI Styles
+CUSTOM_UI = """
+<style>
+    :root {
+        /* Main colors (lighter versions) */
+        --primary-teal: rgb(67, 163, 151);          
+        --primary-teal-light: rgba(67, 163, 151, 0.1);
+        --neutral-gray: rgb(220, 220, 220);         
+        --primary-blue: rgb(82, 142, 198);          
+        --primary-blue-light: rgba(82, 142, 198, 0.1);
+        --dark-blue: rgb(65, 105, 145);             
+        --cream-bg: rgba(218,6,18, 0.05);             
+        --accent-red: rgb(235, 108, 108);           
+        --accent-red-light: rgba(235, 108, 108, 0.1);
+        --white: #ffffff;
+        --text-dark: rgb(60, 60, 60);
+    }
+
+    /* Global Styles */
+    .stApp {
+        background-color: var(--white);
+        color: var(--text-dark);
+        font-family: "Roboto", sans-serif;
+    }
+
+    /* Headers */
+    h1, h2, h3 {
+        color: var(--dark-blue);
+        border-bottom: 2px solid var(--primary-teal-light);
+        padding-bottom: 0.2em;
+        background: linear-gradient(to right, var(--white) 0%, var(--primary-teal-light) 100%);
+        background-clip: text;
+        -webkit-background-clip: text;
+        color: var(--dark-blue); /* re-apply color since gradient would show through transparent text */
+        text-shadow: 0 1px 1px rgba(0,0,0,0.1);
+        font-weight: 600;
+    }
+
+    /* Chat Messages */
+    .stChatMessage {
+        background: var(--primary-blue-light);
+        border-left: 3px solid var(--primary-blue);
+        border-radius: 6px;
+        box-shadow: 0 1px 3px rgba(0,0,0,0.1);
+        padding: 0.5em;
+        transition: box-shadow 0.2s ease;
+    }
+    .stChatMessage:hover {
+        box-shadow: 0 2px 5px rgba(0,0,0,0.2);
+    }
+
+    /* Buttons */
+    .stButton > button {
+        background: var(--white);
+        color: var(--primary-teal);
+        border: 1px solid var(--primary-teal);
+        border-radius: 4px;
+        transition: background 0.2s ease, color 0.2s ease, box-shadow 0.2s ease;
+        font-weight: 500;
+    }
+    .stButton > button:hover {
+        background: var(--primary-teal);
+        color: var(--white);
+        box-shadow: 0 2px 5px rgba(0,0,0,0.15);
+    }
+
+    /* Fixed Button (if any) */
+    .fixed-button button {
+        background: var(--primary-blue);
+        color: var(--white);
+        border-radius: 4px;
+    }
+
+    /* Sidebar */
+    [data-testid=stSidebar] {
+        background-color: var(--cream-bg);
+        color: var(--text-dark);
+        border-right: 1px solid var(--neutral-gray);
+    }
+    [data-testid=stSidebar] .stSelectbox label {
+        color: var(--text-dark);
+    }
+
+    /* Input Fields */
+    .stTextInput > div > div > input {
+        background: var(--white);
+        border: 1px solid var(--neutral-gray);
+        border-radius: 4px;
+        padding: 0.4em;
+        transition: border-color 0.2s ease, box-shadow 0.2s ease;
+    }
+    .stTextInput > div > div > input:focus {
+        border-color: var(--primary-blue);
+        box-shadow: 0 0 0 1px var(--primary-blue-light);
+    }
+
+    /* Alerts and Messages */
+    .stAlert {
+        background: var(--primary-teal-light);
+        border-left: 3px solid var(--primary-teal);
+        border-radius: 4px;
+        padding: 0.5em;
+        box-shadow: 0 1px 3px rgba(0,0,0,0.1);
+    }
+    .stAlert.error {
+        background: var(--accent-red-light);
+        border-left: 3px solid var(--accent-red);
+    }
+
+    /* DataFrame */
+    .stDataFrame {
+        border: 1px solid var(--neutral-gray);
+        border-radius: 4px;
+        overflow: hidden;
+        box-shadow: 0 1px 3px rgba(0,0,0,0.1);
+    }
+
+    /* Expander */
+    .streamlit-expanderHeader {
+        background: var(--white);
+        border: 1px solid var(--neutral-gray);
+        border-radius: 4px;
+        font-weight: 500;
+    }
+    .streamlit-expanderHeader:hover {
+        box-shadow: 0 1px 3px rgba(0,0,0,0.1);
+    }
+
+    /* Checkbox */
+    .stCheckbox > label > div[role="checkbox"] {
+        border-color: var(--primary-teal);
+        border-radius: 3px;
+        transition: background 0.2s ease;
+    }
+    .stCheckbox > label > div[role="checkbox"][aria-checked="true"] {
+        background-color: var(--primary-teal);
+    }
+
+    /* Search bar */
+    .stSearchInput > div > div > input {
+        border: 1px solid var(--primary-blue);
+        border-radius: 4px;
+    }
+    .stSearchInput > div > div > input:focus {
+        border-color: var(--primary-teal);
+        box-shadow: 0 0 0 1px var(--primary-teal-light);
+    }
+
+    /* Secondary buttons */
+    [data-testid="stButton"] > button[kind="secondary"] {
+        background: var(--white);
+        color: var(--primary-blue);
+        border: 1px solid var(--primary-blue);
+        border-radius: 4px;
+        transition: background 0.2s ease, color 0.2s ease, box-shadow 0.2s ease;
+    }
+    [data-testid="stButton"] > button[kind="secondary"]:hover {
+        background: var(--primary-blue);
+        color: var(--white);
+        box-shadow: 0 2px 5px rgba(0,0,0,0.15);
+    }
+
+    /* Delete or cancel actions */
+    .delete-button > button {
+        color: var(--accent-red);
+        border-color: var(--accent-red);
+        border-radius: 4px;
+        transition: background 0.2s ease, color 0.2s ease;
+    }
+    .delete-button > button:hover {
+        background: var(--accent-red);
+        color: var(--white);
+    }
+</style>
+"""
+
+
+
+# Constants
+SYSTEM_ICON = "img/11111111.png"
+USER_ICON = "img/2222222.png"

src/utils.py

@@ -0,0 +1,64 @@
+# src/utils.py
+
+
+import os
+import uuid
+import re
+import logging
+import pandas as pd
+import matplotlib.pyplot as plt
+import streamlit as st
+import time
+import json
+
+# Generate a unique image path for saving plots
+def generate_unique_image_path():
+    figs_dir = os.path.join('tmp', 'figs')
+    os.makedirs(figs_dir, exist_ok=True)
+    unique_filename = f'fig_{uuid.uuid4()}.png'
+    unique_path = os.path.join(figs_dir, unique_filename)
+    logging.debug(f"Generated unique image path: {unique_path}")
+    return unique_path
+
+
+# Function to sanitize input
+def sanitize_input(query: str) -> str:
+    return query.strip()
+
+# Define the function to extract the last Python REPL command
+def get_last_python_repl_command():
+    import streamlit as st  # Ensure Streamlit is imported
+    if 'intermediate_steps' not in st.session_state:
+        logging.warning("No intermediate steps found in session state.")
+        return None
+
+    intermediate_steps = st.session_state['intermediate_steps']
+    python_repl_commands = []
+    for step in intermediate_steps:
+        action = step[0]
+        observation = step[1]
+        if action.get('tool') == 'Python_REPL':
+            python_repl_commands.append(action)
+
+    if python_repl_commands:
+        last_command_action = python_repl_commands[-1]
+        command = last_command_action.get('tool_input', '')
+        logging.debug(f"Extracted last Python REPL command: {command}")
+        return command
+    else:
+        logging.warning("No Python_REPL commands found in intermediate steps.")
+        return None
+
+
+def log_history_event(session_data: dict, event_type: str, details: dict):
+    if "execution_history" not in session_data:
+        session_data["execution_history"] = []  # fallback
+
+    timestamp = time.strftime("%Y-%m-%d %H:%M:%S", time.localtime())
+    event = {
+        "type": event_type,
+        "timestamp": timestamp
+    }
+    event.update(details)  # merges in content from details
+
+    session_data["execution_history"].append(event)