104 lines
3.2 KiB
Python
104 lines
3.2 KiB
Python
import asyncio
|
|
import time
|
|
|
|
import uvicorn
|
|
from fastapi import FastAPI, Security, HTTPException
|
|
from fastapi.security import APIKeyHeader
|
|
from pydantic import BaseModel
|
|
|
|
from utils.langgraph_agent import RULES
|
|
from utils.secrets_manager import SECRETS
|
|
from services.document_extractor import parse_s3_uri, extract_text_from_s3_document
|
|
from services.authorization import evaluate_servico
|
|
from services.result_store import save_results
|
|
|
|
app = FastAPI()
|
|
|
|
_api_key_header = APIKeyHeader(name="X-API-Key")
|
|
API_KEY = SECRETS["API-KEY"]
|
|
|
|
|
|
def verify_api_key(api_key: str = Security(_api_key_header)):
|
|
if api_key != API_KEY:
|
|
raise HTTPException(status_code=403, detail="Invalid API key")
|
|
return api_key
|
|
|
|
|
|
async def process_guia(guia: dict) -> dict:
|
|
t_start = time.time()
|
|
anexos = guia.get("anexos", [])
|
|
|
|
async def _extract_anexo(anexo_idx: int, anexo: dict):
|
|
s3_uri = anexo.get("urlAnexo") or anexo.get("URLAnexo", "")
|
|
nome_arquivo = anexo.get("nomeArquivo", f"attachment_{anexo_idx}")
|
|
if not s3_uri or not s3_uri.startswith("s3://"):
|
|
anexo["textoExtraido"] = ""
|
|
return None
|
|
t0 = time.time()
|
|
try:
|
|
bucket, key = parse_s3_uri(s3_uri)
|
|
extracted_text, page_count = await extract_text_from_s3_document(bucket, key)
|
|
except Exception as e:
|
|
extracted_text = ""
|
|
page_count = 0
|
|
anexo["error"] = str(e)
|
|
anexo["textoExtraido"] = extracted_text
|
|
anexo["pageCount"] = page_count
|
|
anexo["tempoExtracaoSegundos"] = round(time.time() - t0, 2)
|
|
return f"--- {nome_arquivo} ---\n{extracted_text}"
|
|
|
|
t_extracao_start = time.time()
|
|
parts = await asyncio.gather(*[_extract_anexo(i, a) for i, a in enumerate(anexos)])
|
|
file_content = "\n\n".join(p for p in parts if p)
|
|
t_extracao = round(time.time() - t_extracao_start, 2)
|
|
|
|
servicos = guia.get("servicos", [])
|
|
t_agent_start = time.time()
|
|
guia["avaliacaoAgente"] = list(
|
|
await asyncio.gather(*[evaluate_servico(s, guia, file_content) for s in servicos])
|
|
)
|
|
t_agent = round(time.time() - t_agent_start, 2)
|
|
|
|
guia["tempoProcessamento"] = {
|
|
"extracaoSegundos": t_extracao,
|
|
"agentSegundos": t_agent,
|
|
"totalSegundos": round(time.time() - t_start, 2),
|
|
}
|
|
return guia
|
|
|
|
|
|
class ProcessRequest(BaseModel):
|
|
operadora: dict
|
|
guias: list[dict]
|
|
|
|
|
|
@app.post("/process", dependencies=[Security(verify_api_key)])
|
|
async def process(request: ProcessRequest):
|
|
raw_results = await asyncio.gather(
|
|
*[process_guia(guia) for guia in request.guias],
|
|
return_exceptions=True,
|
|
)
|
|
results = [
|
|
{"error": str(r), "guia": request.guias[i].get("guia", {}).get("codigoGuiaLocal", f"index_{i}")}
|
|
if isinstance(r, Exception) else r
|
|
for i, r in enumerate(raw_results)
|
|
]
|
|
|
|
await save_results(results)
|
|
|
|
return {"status": "success", "operadora": request.operadora, "guias": results}
|
|
|
|
|
|
@app.get("/health")
|
|
async def health():
|
|
return {"status": "healthy"}
|
|
|
|
|
|
@app.get("/rules")
|
|
async def get_rules():
|
|
return {"codes": list(RULES.keys())}
|
|
|
|
|
|
if __name__ == "__main__":
|
|
uvicorn.run(app, host="0.0.0.0", port=8000)
|