""" 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()