vent_ai_platform/app/agents/context.py

from __future__ import annotations

from sqlalchemy.orm import Session

from app.db.repositories import diagnosis as diagnosis_repo
from app.db.repositories import document as document_repo
from app.db.repositories import evaluation as evaluation_repo
from app.db.repositories import knowledge as knowledge_repo
from app.db.repositories import monitoring as monitoring_repo
from app.db.repositories import work_order as work_order_repo
from app.services.mine_service import resolve_mine


def build_mine_context(
    db: Session,
    mine_id: str | None,
    *,
    include_documents: bool = False,
) -> dict:
    mine = resolve_mine(db, mine_id)
    context = {
        "mine": {
            "id": mine.id,
            "code": mine.code,
            "name": mine.name,
            "ventilation_method": mine.ventilation_method,
        },
        "monitoring": monitoring_repo.monitoring_overview(db, mine.id),
        "diagnosis": diagnosis_repo.result_summary(db, mine.id),
        "work_orders": work_order_repo.status_summary(db, mine.id),
    }
    runs = evaluation_repo.list_runs(db, mine.id, limit=1)
    if runs:
        context["latest_evaluation"] = {
            "run_no": runs[0].run_no,
            "score": runs[0].total_score,
            "grade": runs[0].grade,
            "summary": runs[0].summary,
        }
    if include_documents:
        documents, total = document_repo.list_documents(
            db,
            mine_id=mine.id,
            page=1,
            page_size=10,
        )
        context["documents"] = {
            "total": total,
            "recent": [
                {
                    "id": item.id,
                    "title": item.title,
                    "category": item.category,
                    "parse_status": item.parse_status,
                }
                for item in documents
            ],
        }
    return context


def build_knowledge_context(
    db: Session,
    mine_id: str | None,
    question: str,
    limit: int = 8,
) -> dict:
    items, total = knowledge_repo.list_knowledge(
        db,
        mine_id=mine_id,
        keyword=question,
        page=1,
        page_size=limit,
    )
    return {
        "total": total,
        "items": [
            {
                "id": item.id,
                "code": item.code,
                "title": item.title,
                "type": item.knowledge_type,
                "category": item.category,
                "summary": item.summary,
                "content": item.content[:1200],
            }
            for item in items
        ],
    }