query_tools/logbook_availability_query.py

import pandas as pd
import os
from datetime import datetime


# ============================
# 自动定位项目路径
# ============================

BASE_DIR = os.path.dirname(
    os.path.dirname(
        os.path.abspath(__file__)
    )
)


DATA_PATH = os.path.join(
    BASE_DIR,
    "data",
    "logbook_availability.csv"
)


OUTPUT_DIR = os.path.join(
    BASE_DIR,
    "output"
)



# ============================
# markdown生成函数
# 替代 pandas.to_markdown()
# ============================

def dataframe_to_markdown(df, max_rows=10):

    if df.empty:
        return "暂无数据"


    df=df.head(max_rows)


    columns=df.columns.tolist()


    md="| " + " | ".join(columns)+" |\n"

    md+="| "+" | ".join(
        ["---"]*len(columns)
    )+" |\n"


    for _,row in df.iterrows():

        md+="| "+" | ".join(
            str(x)
            for x in row.tolist()
        )+" |\n"


    return md



# ============================
# 参数检查
# ============================

def check_year(value,name):

    if value is None:
        return None


    try:
        return int(value)

    except:

        raise ValueError(
            f"{name}必须是年份数字，例如2020"
        )




# ============================
# 主查询函数
# ============================

def query_logbook_availability(
        region=None,
        year_start=None,
        year_end=None,
        species=None,
        data_type=None,
        output_format="markdown"
):

    """
    查询logbook数据可用性

    返回:
    summary
    records
    preview_markdown
    csv_path
    excel_path
    source_files

    """


    try:

        # ------------------
        # 参数校验
        # ------------------

        year_start = check_year(
            year_start,
            "year_start"
        )

        year_end = check_year(
            year_end,
            "year_end"
        )



        if year_start and year_end:

            if year_start > year_end:

                return {

                    "summary":
                    "错误：开始年份不能大于结束年份",

                    "records":[]

                }



        # ------------------
        # 文件检查
        # ------------------

        if not os.path.exists(DATA_PATH):

            return {

                "summary":
                f"找不到数据文件:{DATA_PATH}",

                "records":[]

            }



        # ------------------
        # 读取数据
        # ------------------

        df=pd.read_csv(
            DATA_PATH
        )


        # year强制转换

        if "year" in df.columns:

            df["year"]=pd.to_numeric(
                df["year"],
                errors="coerce"
            )


        result=df.copy()



        # ------------------
        # 条件过滤
        # ------------------

        if region:

            result=result[
                result["region"]
                .astype(str)
                .str.contains(
                    region,
                    na=False
                )
            ]



        if species:

            result=result[
                result["species"]
                .astype(str)
                .str.contains(
                    species,
                    na=False
                )
            ]



        if data_type:

            result=result[
                result["data_type"]
                .astype(str)
                .str.contains(
                    data_type,
                    na=False
                )
            ]



        if year_start:

            result=result[
                result.year>=year_start
            ]



        if year_end:

            result=result[
                result.year<=year_end
            ]



        # ------------------
        # 无结果
        # ------------------

        if result.empty:

            return {

                "summary":
                "没有找到符合条件的logbook数据",

                "records":[],

                "preview_markdown":
                "暂无数据",

                "source_files":
                [
                    DATA_PATH
                ]

            }



        # ------------------
        # 统计信息
        # ------------------

        summary={

            "records_count":
            len(result),

            "year_range":
            [
                int(result.year.min()),
                int(result.year.max())
            ],

            "regions":
            result.region.unique().tolist(),

            "species":
            result.species.unique().tolist()

        }



        # ------------------
        # 输出文件
        # ------------------

        os.makedirs(
            OUTPUT_DIR,
            exist_ok=True
        )


        timestamp=datetime.now()\
            .strftime("%Y%m%d_%H%M%S")


        csv_path=None
        excel_path=None



        if output_format in [
            "csv",
            "excel"
        ]:

            csv_path=os.path.join(
                OUTPUT_DIR,
                f"logbook_{timestamp}.csv"
            )


            result.to_csv(
                csv_path,
                index=False,
                encoding="utf-8-sig"
            )



        if output_format=="excel":

            excel_path=os.path.join(
                OUTPUT_DIR,
                f"logbook_{timestamp}.xlsx"
            )


            result.to_excel(
                excel_path,
                index=False
            )



        # ------------------
        # Agent标准返回
        # ------------------

        return {


            "summary":
            summary,


            "records":
            result.to_dict(
                orient="records"
            ),


            "preview_markdown":
            dataframe_to_markdown(
                result
            ),


            "csv_path":
            csv_path,


            "excel_path":
            excel_path,


            "source_files":
            [
                DATA_PATH
            ]

        }



    except Exception as e:


        return {

            "summary":
            f"查询失败:{str(e)}",

            "records":[]

        }