server.py

"""
lex-mcp — Assistente Jurídico via MCP + Hugging Face
Expõe ferramentas especializadas em direito para qualquer LLM host compatível com MCP.
"""
from __future__ import annotations  # ← CORREÇÃO: adicionar __
import os
import re
from typing import Any
import httpx
from fastmcp import FastMCP
from huggingface_hub import HfApi, list_models  # ← CORREÇÃO: remover ModelFilter
from huggingface_hub.utils import RepositoryNotFoundError
from datasets import load_dataset, get_dataset_config_names, get_dataset_split_names

# ── Bootstrap ─────────────────────────────────────────────────────────────────
mcp = FastMCP(
    name="lex-mcp",
    instructions="""
Você é LEX, um assistente jurídico especializado alimentado por modelos e datasets
do Hugging Face Hub. Você possui quatro ferramentas:
• search_legal_models   — encontra modelos de NLP treinados em domínio jurídico
• explore_legal_dataset — inspeciona datasets jurídicos (jurisprudência, leis, contratos)
• analyze_legal_text    — roda inferência NLP em texto jurídico (classificação, NER, resumo)
• find_jurisprudence    — busca decisões e ementas em datasets de jurisprudência

IMPORTANTE: Sempre use as ferramentas para buscar dados atuais do Hub.
Nunca invente modelos ou citações. Indique limitações quando relevante.
Responda em português quando o usuário escrever em português.
""",
)

HF_TOKEN = os.getenv("HF_TOKEN")
api = HfApi(token=HF_TOKEN)

# Modelos jurídicos de referência no HF Hub (curados)
LEGAL_MODEL_HINTS = [
    "legal", "juridico", "jurídico", "law", "legislation",
    "bert-legal", "legalbert", "law-bert", "contracts", "court",
    "nlp-laval", "legal-xlm", "legalbench", "saul", "brazilianLegal",
]

LEGAL_DATASET_HINTS = [
    "legal", "law", "court", "jurisprudence", "legislation",
    "contracts", "case-law", "oab", "stf", "stj", "tjsp",
]

# ── Tool 1 — search_legal_models ──────────────────────────────────────────────
@mcp.tool(
    description=(
        "Busca modelos de NLP especializados em domínio jurídico no Hugging Face Hub.  "
        "Filtre por língua (ex: 'pt' para português), tarefa (ex: 'text-classification',  "
        "'token-classification', 'summarization') e palavras-chave.  "
        "Retorna os modelos mais baixados com metadados completos. "
    )
)
def search_legal_models(
    query: str = "legal",
    language: str = "pt",
    task: str = "",
    limit: int = 8,
) -> list[dict[str, Any]]:
    """Retorna modelos jurídicos ordenados por downloads."""
    # Enriquecer query com termos jurídicos se necessário
    legal_query = query if any(h in query.lower() for h in LEGAL_MODEL_HINTS) else f"legal {query}"

    # ← CORREÇÃO: Usar parâmetros diretos ao invés de ModelFilter
    results = list(
        list_models(
            search=legal_query,
            language=language or None,
            pipeline_tag=task or None,  # ← CORREÇÃO: pipeline_tag ao invés de task
            sort="downloads",
            direction=-1,
            limit=limit,
            token=HF_TOKEN,
            cardData=True,
        )
    )

    return [
        {
            "id": m.modelId,
            "task": m.pipeline_tag,
            "downloads": m.downloads,
            "likes": m.likes,
            "last_modified": str(m.lastModified)[:10],
            "tags": [t for t in (m.tags or []) if len(t) < 40][:8],
            "language": getattr(m, "language", None),
            "hf_url": f"https://huggingface.co/{m.modelId}",
        }
        for m in results
    ]

# ── Tool 2 — explore_legal_dataset ───────────────────────────────────────────
@mcp.tool(
    description=(
        "Inspeciona um dataset jurídico no Hugging Face Hub.  "
        "Retorna configs disponíveis, splits, schema de colunas e exemplos de registros.  "
        "Ideal para entender datasets de jurisprudência, legislação e contratos.  "
        "Use dataset_id como 'joelniklaus/MultiLegalPile', 'lexlms/lex_glue', etc. "
    )
)
def explore_legal_dataset(
    dataset_id: str,
    config: str = "default",
    split: str = "train",
    n_samples: int = 3,
) -> dict[str, Any]:
    """Retorna schema + amostras de um dataset jurídico."""
    try:
        configs = get_dataset_config_names(dataset_id, token=HF_TOKEN)
    except Exception:
        configs = [config]

    resolved_config = config if config in configs else (configs[0] if configs else None)

    try:
        splits = get_dataset_split_names(dataset_id, config_name=resolved_config, token=HF_TOKEN)
    except Exception:
        splits = [split]

    # ← CORREÇÃO: remover espaço em "spli t"
    resolved_split = split if split in splits else (splits[0] if splits else "train")

    try:
        ds = load_dataset(
            dataset_id,
            name=resolved_config,
            split=f"{resolved_split}[:{n_samples}]",
            token=HF_TOKEN,
            trust_remote_code=False,
        )
        features = {k: str(v) for k, v in ds.features.items()}
        samples = ds.to_list()
        # Truncar textos longos para não explodir o contexto
        for sample in samples:
            for key, val in sample.items():
                if isinstance(val, str) and len(val) > 600:
                    sample[key] = val[:600] + "…"
    except Exception as e:
        features = {}
        samples = []
        return {
            "dataset_id": dataset_id,
            "error": str(e),
            "configs_available": configs,
            "splits_available": splits,
        }

    return {
        "dataset_id": dataset_id,
        "hf_url": f"https://huggingface.co/datasets/{dataset_id}",
        "configs_available": configs,
        "splits_available": splits,
        "resolved": {"config": resolved_config, "split": resolved_split},
        "total_features": len(features),
        "features": features,
        "samples": samples,
    }

# ── Tool 3 — analyze_legal_text ──────────────────────────────────────────────
@mcp.tool(
    description=(
        "Roda inferência NLP em texto jurídico usando a Hugging Face Inference API.  "
        "Tarefas suportadas: 'summarization' (resumo de decisões),  "
        "'text-classification' (classificação de matéria/área do direito),  "
        "'token-classification' (NER: partes, datas, valores),  "
        "'question-answering' (responde perguntas sobre o texto).  "
        "Se model_id não for fornecido, usa modelos jurídicos recomendados. "
    )
)
def analyze_legal_text(
    text: str,
    task: str = "summarization",
    model_id: str = "",
    context: str = "",
) -> dict[str, Any]:
    """Executa análise NLP jurídica via Inference API."""
    # Modelos padrão por tarefa (jurídicos ou multilíngues de qualidade)
    DEFAULT_MODELS: dict[str, str] = {
        "summarization": "facebook/bart-large-cnn",
        "text-classification": "nlpaueb/legal-bert-base-uncased",
        "token-classification": "nlpaueb/legal-bert-base-uncased",
        "question-answering": "deepset/roberta-base-squad2",
        "fill-mask": "nlpaueb/legal-bert-base-uncased",
    }

    resolved_model = model_id or DEFAULT_MODELS.get(task, "facebook/bart-large-cnn")

    url = f"https://api-inference.huggingface.co/models/{resolved_model}"
    headers = {"Content-Type": "application/json"}
    if HF_TOKEN:
        headers["Authorization"] = f"Bearer {HF_TOKEN}"
    if task:
        headers["X-Task"] = task

    # Montar payload conforme a tarefa
    if task == "question-answering" and context:
        payload: dict[str, Any] = {"inputs": {"question": text, "context": context}}
    elif task == "summarization":
        # Truncar para evitar erros de tamanho máximo
        payload = {
            "inputs": text[:1024],
            "parameters": {"max_length": 200, "min_length": 40, "do_sample": False},
        }
    else:
        payload = {"inputs": text[:512]}

    with httpx.Client(timeout=60.0) as client:
        resp = client.post(url, headers=headers, json=payload)

    if resp.status_code == 503:
        return {
            "status": "model_loading",
            "model_id": resolved_model,
            "message": "Modelo está carregando. Tente novamente em 20-30 segundos.",
        }

    if resp.status_code != 200:
        return {
            "error": f"HTTP {resp.status_code}",
            "model_id": resolved_model,
            "detail": resp.text[:400],
        }

    try:
        result = resp.json()
    except Exception:
        result = resp.text

    return {
        "model_id": resolved_model,
        "task": task,
        "hf_url": f"https://huggingface.co/{resolved_model}",
        "result": result,
    }

# ── Tool 4 — find_jurisprudence ───────────────────────────────────────────────
@mcp.tool(
    description=(
        "Busca decisões judiciais e ementas em datasets de jurisprudência disponíveis  "
        "no Hugging Face Hub. Pesquisa por palavras-chave no texto das decisões.  "
        "Retorna ementas, tribunal, data e número do processo quando disponíveis.  "
        "Datasets suportados: 'joelniklaus/brazilian_court_decisions',  "
        "'lagepaul/jurisprudencia-brasil' e outros datasets jurídicos brasileiros. "
    )
)
def find_jurisprudence(
    keywords: str,
    dataset_id: str = "joelniklaus/brazilian_court_decisions",
    max_results: int = 5,
    split: str = "train",
) -> dict[str, Any]:
    """Busca decisões judiciais por palavras-chave."""
    try:
        configs = get_dataset_config_names(dataset_id, token=HF_TOKEN)
        resolved_config = configs[0] if configs else None
    except Exception:
        resolved_config = None

    try:
        # Carregar slice generoso para fazer busca textual
        ds = load_dataset(
            dataset_id,
            name=resolved_config,
            split=f"{split}[:500]",
            token=HF_TOKEN,
            trust_remote_code=False,
        )
    except Exception as e:
        return {"error": str(e), "dataset_id": dataset_id}

    # Identificar coluna de texto principal
    text_cols = [
        col for col in ds.column_names
        if any(kw in col.lower() for kw in ["text", "ementa", "decision", "body", "content", "acordao"])
    ]
    text_col = text_cols[0] if text_cols else ds.column_names[0]

    # Busca por keywords (case-insensitive)
    kw_pattern = re.compile("|".join(re.escape(k.strip()) for k in keywords.split(",")), re.IGNORECASE)

    matches = []
    for row in ds:
        haystack = str(row.get(text_col, ""))
        if kw_pattern.search(haystack):
            snippet = haystack[:500] + ("…" if len(haystack) > 500 else "")
            matches.append({
                "snippet": snippet,
                "columns": {k: str(v)[:200] for k, v in row.items() if k != text_col},
            })
        if len(matches) >= max_results:
            break

    return {
        "dataset_id": dataset_id,
        "hf_url": f"https://huggingface.co/datasets/{dataset_id}",
        "keywords_searched": keywords,
        "text_column_used": text_col,
        "total_matches": len(matches),
        "results": matches,
    }

# ── Entry point ───────────────────────────────────────────────────────────────
# ← CORREÇÃO: adicionar __
if __name__ == "__main__":
    mcp.run()