main.py - master - Vladimir_Py/cp_generator

main.py
206 строк8 KB
Test Bot
docker ready
05 май 2026, 00:15
05 май 2026, 00:156eecf25
import logging
import os
from pathlib import Path

import pandas as pd
from crewai import CrewOutput
from crewai.types.streaming import CrewStreamingOutput

from schemas.proj_classes import CategoriesSchema
from src.agents.agents_pipeline import run_commercial_proposal_crew
from src.agents.eval_pipeline import run_commercial_proposal_eval
from src.file_utils import read_json, read_csv, save_csv, read_txt, save_txt
from src.llm_factory import LLMFactory

from src.classifier import ProjectClassifier

# Настраиваем логгер
logging.basicConfig(level=logging.DEBUG, format='%(asctime)s | %(levelname)-8s | %(name)s | %(message)s')
logger = logging.getLogger("pipeline")

try:
    # 1. Загружаем сырой JSON и сразу валидируем его через Pydantic!
    raw_json = read_json("src/ai_classes.json")
    categories_schema = CategoriesSchema.model_validate(raw_json)

    # 2. Создаем LLM (берем легкую и дешевую модель для классификации)
    llm = LLMFactory.from_config("openai", tier="lite")

    # 3. Инициализируем классификатор
    pr_classifier = ProjectClassifier(llm=llm, schema=categories_schema)
except Exception as e:
    logger.error(f"Ошибка при инициализации: {e}")
    raise

def single_classify(tz_text : str, classifier : ProjectClassifier):

    # Классифицируем текст
    tz_class = classifier.classify(tz_text)

    # вывод результата
    if tz_class.is_target:
        logger.info(f"✅ В работу! Запускаем пайплайн для: {tz_class.category_key}")
    else:
        logger.info(f"❌ Отбраковано. Код исключения: {tz_class.category_key}")

    return tz_class

def csv_classify(
        csv_path: str,
        classifier: ProjectClassifier,
        column_name: str = "description_clean",
        limit: int = 5,
        output_csv: str = "output.csv"
) -> pd.DataFrame | None:

    logger.info(f"Начинаем пакетный прогон для: {csv_path}")
    df_new = pd.DataFrame()

    # 1. ЗАГРУЗКА ИСТОРИИ (КЭША)
    processed_ids = set()
    if os.path.exists(output_csv):
        try:
            # Читаем существующий лог успехов
            existing_df = pd.read_csv(output_csv)
            if 'project_id' in existing_df.columns:
                # Приводим к строке, чтобы избежать проблем с типами (int vs str)
                processed_ids = set(existing_df['project_id'].astype(str))
            logger.info(f"В кэше найдено {len(processed_ids)} уже обработанных проектов.")
        except Exception as e:
            logger.warning(f"Не удалось прочитать историю из {output_csv}: {e}")

    processed_data = []

    try:
        orders = read_csv(csv_path) # возвращает список словарей
        processed_count = 0

        for idx, order in enumerate(orders, 1):
            # Останавливаемся, если обработали нужное количество НОВЫХ заказов
            if processed_count >= limit:
                logger.info(f"Достигнут лимит в {limit} новых обработок.")
                break

            project_id = str(order.get("project_id", "")).strip()

            # 2. ПРОВЕРКА НА ДУБЛИКАТ
            if project_id and project_id in processed_ids:
                logger.debug(f"[{idx}] Пропуск: проект {project_id} уже обработан.")
                continue

            text = order.get(column_name, "").strip()
            if not text:
                continue

            # 3. ЗАПУСК ИИ (Только для новых)
            res = classifier.classify(text)
            processed_count += 1

            status_icon = "✅" if res.is_target else "❌"
            logger.debug(f"[Новый #{processed_count}] {status_icon} {res.category_key:<18} | {text[:80]}...")

            processed_data.append({
                "project_id": project_id,
                "description_clean": order.get("description_clean"),
                "budget_max": order.get("budget_max"),
                "timeline_max_days": order.get("timeline_max_days"),
                "is_target": res.is_target,
                "category": res.category_key
            })

        df_new = pd.DataFrame(processed_data)
        logger.info(f"Прогон завершен. Новых обработано: {df_new.shape[0]}.")

        return df_new

    except FileNotFoundError as e:
        logger.error(f"Ошибка: файл не найден: {e}")
    except Exception as e:
        logger.error(f"Ошибка при пакетной обработке: {e}")
        # Возвращаем DataFrame с результатами обработки до ошибки или пустой
        return df_new

def create_cp(category: str, tz_source: str) -> CrewOutput | CrewStreamingOutput:
    """
    Точка входа для генерации КП по ТЗ.
    tz_source- пусть к файлу (например, 'task_example.md').
    """
    try:
        # 1. Читаем ТЗ из файла
        if os.path.isfile(tz_source):
            logger.info(f"Читаем ТЗ из файла: {tz_source}")
            with open(tz_source, 'r', encoding='utf-8') as f:
                tz_text = f.read()
        else:
            logger.error(f"Ошибка: файл не найден: {tz_source}")
            raise FileNotFoundError

        # 2. Инициализируем  модель
        llm_agents = LLMFactory.from_config("gemini", tier="lite",
                                            config_path="llm_config.yaml", for_crewai = True)

        # 3. Запускаем конвейер
        logger.info(f"Передаем ТЗ в пайплайн (категория: {category})...")
        final_cp = run_commercial_proposal_crew(category, tz_text, llm_agents,
                                                agent_config_path='src/agents/agents_config.json')

        return final_cp

    except Exception as e:
        logger.error(f"Ошибка при оценке ТЗ: {e}")
        return CrewOutput()

def evaluate_cp(tz_source: str, cp_source: str) -> dict:
    """
    Точка входа для оценки КП.
    tz_source- пусть к файлу с исходным ТЗ (например, 'task_example.md')
    cp_source- пусть к файлу КП
    """
    try:
        # 1. Читаем ТЗ и КП из файла

        logger.info(f"Читаем ТЗ из файла: {tz_source}")
        tz_text = read_txt(tz_source)
        logger.info(f"Читаем КП из файла: {cp_source}")
        cp_text = read_txt(cp_source)


        # 2. Инициализируем  модель
        llm_agents = LLMFactory.from_config("openrouter", tier="full",
                                            config_path="llm_config.yaml", for_crewai = True)

        # 3. Запускаем конвейер
        logger.info(f"Передаем ТЗ и КП в пайплайн ...")
        final_cp = run_commercial_proposal_eval( tz_text, cp_text, llm_agents,
                                                agent_config_path='src/agents/agents_config.json')

        return final_cp

    except Exception as e:
        logger.error(f"Ошибка при оценке ТЗ: {e}")
        return {}


if __name__ == "__main__":
    # csv_in_path = "_data/input_tz/fl_orders.csv"
    # csv_out_path = "_data/classifications/classified_orders_1_gpt4.csv"
    # classified_df = csv_classify(csv_path=csv_in_path, classifier=pr_classifier, limit=200,
    #                              output_csv = csv_out_path)
    # save_csv(data = classified_df, filepath=csv_out_path)

    # target_category = "other_ai_ml"
    # tz_path = "src/agents/examples"
    # tz_cat = "video_analytics"
    # tz_file =  "task_example.md"
    # tz_fullpath = Path (tz_path) /tz_cat/ tz_file
    #
    # result = create_cp(target_category, tz_fullpath)
    # print(result)
    # cp_path = Path ("_data/output_kp")
    # cp_file = tz_cat + "_cp.md"
    # cp_fullpath = Path (cp_path) / cp_file
    # save_txt(text=result.__str__(), filepath=cp_fullpath)


    result = evaluate_cp("src/agents/examples/video_analytics/task_example.md",
                         "_data/output_kp/writer_example.md")