Создание DataFrame-агента на базе Gemini для анализа данных на естественном языке с Pandas и LangChain

!pip install langchain_experimental langchain_google_genai pandas

import os
import pandas as pd
import numpy as np
from langchain.agents.agent_types import AgentType
from langchain_experimental.agents.agent_toolkits import create_pandas_dataframe_agent
from langchain_google_genai import ChatGoogleGenerativeAI
 
os.environ["GOOGLE_API_KEY"] = "Use Your Own API Key"

def setup_gemini_agent(df, temperature=0, model="gemini-1.5-flash"):
    llm = ChatGoogleGenerativeAI(
        model=model,
        temperature=temperature,
        convert_system_message_to_human=True
    )
   
    agent = create_pandas_dataframe_agent(
        llm=llm,
        df=df,
        verbose=True,
        agent_type=AgentType.OPENAI_FUNCTIONS,
        allow_dangerous_code=True
    )
    return agent

def load_and_explore_data():
    print("Загрузка набора данных Titanic...")
    df = pd.read_csv(
        "https://raw.githubusercontent.com/pandas-dev/pandas/main/doc/data/titanic.csv"
    )
    print(f"Размер набора данных: {df.shape}")
    print(f"Столбцы: {list(df.columns)}")
    return df

def basic_analysis_demo(agent):
    print("\nДЕМОНСТРАЦИЯ БАЗОВОГО АНАЛИЗА")
    print("=" * 50)
   
    queries = [
        "Сколько строк и столбцов в наборе данных?",
        "Каков процент выживших?",
        "Сколько людей имеют более 3 братьев или сестер?",
        "Каков квадратный корень среднего возраста?",
        "Покажи распределение пассажирских классов"
    ]
   
    for query in queries:
        print(f"\nЗапрос: {query}")
        try:
            result = agent.invoke(query)
            print(f"Результат: {result['output']}")
        except Exception as e:
            print(f"Ошибка: {e}")

def advanced_analysis_demo(agent):
    print("\nДЕМОНСТРАЦИЯ ПРОДВИНУТОГО АНАЛИЗА")
    print("=" * 50)
   
    advanced_queries = [
        "Какова корреляция между возрастом и стоимостью билета?",
        "Создай анализ выживаемости по полу и классу",
        "Каков медианный возраст для каждого пассажирского класса?",
        "Найди пассажиров с самыми высокими тарифами и их детали",
        "Вычисли процент выживших для разных возрастных групп (0-18, 18-65, 65+)"
    ]
   
    for query in advanced_queries:
        print(f"\nЗапрос: {query}")
        try:
            result = agent.invoke(query)
            print(f"Результат: {result['output']}")
        except Exception as e:
            print(f"Ошибка: {e}")

def multi_dataframe_demo():
    print("\nДЕМОНСТРАЦИЯ МНОЖЕСТВЕННЫХ DATAFRAME")
    print("=" * 50)
   
    df = pd.read_csv(
        "https://raw.githubusercontent.com/pandas-dev/pandas/main/doc/data/titanic.csv"
    )
   
    df_filled = df.copy()
    df_filled["Age"] = df_filled["Age"].fillna(df_filled["Age"].mean())
   
    agent = setup_gemini_agent([df, df_filled])
   
    queries = [
        "Сколько строк с разными значениями возраста между двумя наборами?",
        "Сравни средний возраст в обоих наборах",
        "Какой процент значений возраста отсутствовал в оригинальном наборе?",
        "Покажи сводную статистику по возрасту в обоих наборах"
    ]
   
    for query in queries:
        print(f"\nЗапрос: {query}")
        try:
            result = agent.invoke(query)
            print(f"Результат: {result['output']}")
        except Exception as e:
            print(f"Ошибка: {e}")

def custom_analysis_demo(agent):
    print("\nДЕМОНСТРАЦИЯ КАСТОМНОГО АНАЛИЗА")
    print("=" * 50)
   
    custom_queries = [
        "Создай оценку риска для каждого пассажира на основе: Возраста (чем старше, тем выше риск), Пол (мужчины = выше риск), Класс (третий класс = выше риск), Размер семьи (одинокие или большие семьи = выше риск). Затем покажи 10 пассажиров с самым высоким риском, которые выжили",
       
        "Проанализируй информацию о палубе из данных каюты: извлеки букву палубы, покажи выживаемость по палубам, какая палуба имела самый высокий уровень выживаемости?",
       
        "Найди интересные закономерности: выживали ли вместе люди с одинаковыми фамилиями? Какова связь между ценой билета и выживанием? Были ли возрастные группы с 100% выживаемостью?"
    ]
   
    for i, query in enumerate(custom_queries, 1):
        print(f"\nКастомный анализ {i}:")
        print(f"Запрос: {query[:100]}...")
        try:
            result = agent.invoke(query)
            print(f"Результат: {result['output']}")
        except Exception as e:
            print(f"Ошибка: {e}")

def main():
    print("Продвинутый Pandas агент с Gemini")
    print("=" * 60)
   
    if not os.getenv("GOOGLE_API_KEY"):
        print("Внимание: GOOGLE_API_KEY не установлен!")
        print("Пожалуйста, установите ваш Gemini API ключ в переменные окружения.")
        return
   
    try:
        df = load_and_explore_data()
        print("\nНастройка агента Gemini...")
        agent = setup_gemini_agent(df)
       
        basic_analysis_demo(agent)
        advanced_analysis_demo(agent)
        multi_dataframe_demo()
        custom_analysis_demo(agent)
       
        print("\nРуководство успешно завершено!")
       
    except Exception as e:
        print(f"Ошибка: {e}")
        print("Убедитесь, что вы установили все необходимые пакеты и задали API ключ.")
 
 
if __name__ == "__main__":
    main()

df = pd.read_csv("https://raw.githubusercontent.com/pandas-dev/pandas/main/doc/data/titanic.csv")
agent = setup_gemini_agent(df)
 
agent.invoke("Какие факторы сильнее всего влияли на выживаемость?")
agent.invoke("Создай подробный анализ выживаемости по порту посадки")
agent.invoke("Найди интересные аномалии или выбросы в данных")