NEKReport/main_v2.py

587 lines
20 KiB
Python
Raw Permalink Blame History

This file contains ambiguous Unicode characters

This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.

import os
import time
import json
import base64
import hashlib
import logging
from typing import List, Optional, Dict
from pathlib import Path
import uvicorn
import openai
from fastapi import FastAPI, UploadFile, File, HTTPException, Form, Body
from fastapi.responses import JSONResponse, FileResponse, Response
from fastapi.middleware.cors import CORSMiddleware
from fastapi.staticfiles import StaticFiles
from starlette.requests import Request
from starlette.middleware.base import BaseHTTPMiddleware
from pydantic import BaseModel
from prometheus_fastapi_instrumentator import Instrumentator
from urllib.parse import quote
# Импорт нового движка
from cargo_rag_v2 import CargoRAGEngineV2
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)
# ----------------------------- Конфигурация -----------------------------
SHIPPING_TYPES_FILE = os.path.join(os.path.dirname(__file__), "shipping_types.json")
PROCESSED_SHIPPING_TYPES_FILE = os.path.join(
os.path.dirname(__file__), "shipping_types_processed.json"
)
# OpenAI клиент
#openai_api_key = os.getenv("sk-9cfed30eb6894df5b9a3ffb4b1fb956d")
openai_api_key = "sk-9cfed30eb6894df5b9a3ffb4b1fb956d"
if not openai_api_key:
raise RuntimeError("OPENAI_API_KEY is not set")
openai_client = openai.OpenAI(
api_key=openai_api_key,
base_url=os.getenv("OPENAI_BASE_URL", "http://localhost:8090/v1"),
)
# Инициализация движка
rag = CargoRAGEngineV2(openai_client)
# ------------------ Хранилище сессий (писем) --------------------------
# Ключ: session_id, значение: список писем (каждое письмо словарь)
_session_emails: Dict[str, List[Dict]] = {}
# ------------------ Pydantic модели ------------------
class EmailAttachment(BaseModel):
filename: str
size: int
content: Optional[str] = None # base64
class OutlookEmail(BaseModel):
id: str
subject: str
sender: str
senderName: Optional[str] = None
body: str
body_html: Optional[str] = None
receivedTime: Optional[str] = None
to: Optional[str] = None
cc: Optional[str] = None
attachments: List[EmailAttachment] = []
class OutlookEmailsRequest(BaseModel):
emails: List[OutlookEmail]
session_id: Optional[str] = None
class CargoQueryRequest(BaseModel):
query: str
session_id: Optional[str] = None
top_k: int = 10
class CargoQueryResponse(BaseModel):
answer: str
structured_data: dict
sources: list
total_emails_analyzed: int
class AnalyzeCargoRequest(BaseModel):
email_ids: List[str]
class CargoLearningRequest(BaseModel):
structured_data: dict
session_id: Optional[str] = None
context_preview: Optional[str] = None
notes: Optional[str] = None
class ShippingTypeBase(BaseModel):
name: str
criteria: str = ""
keywords: List[str] = []
employee_email: str = ""
confirmation_template: str = ""
info_request_template: str = ""
class ShippingTypeCreate(ShippingTypeBase):
pass
class ShippingTypeUpdate(ShippingTypeBase):
pass
class ShippingType(ShippingTypeBase):
id: int
# ------------------ Middleware и приложение ------------------
app = FastAPI(title="SEPTEM Cargo RAG System V2")
app.mount("/addin", StaticFiles(directory=os.path.join(os.getcwd(), "addin"), html=True), name="addin")
OLD_HOST = "nec.septem.pro"
NEW_HOST = "nec.clients.septem.pro"
@app.middleware("http")
async def redirect_legacy_host(request, call_next):
host = request.headers.get("host", "").split(":")[0].lower()
if host == OLD_HOST:
target_url = f"https://{NEW_HOST}{request.url.path}"
if request.url.query:
target_url = f"{target_url}?{request.url.query}"
return JSONResponse(
status_code=308,
content={"detail": "Permanent Redirect"},
headers={"Location": target_url},
)
return await call_next(request)
class CollapseDuplicatePathSlashesMiddleware(BaseHTTPMiddleware):
async def dispatch(self, request: Request, call_next):
path = request.scope.get("path") or ""
if "//" in path:
collapsed = path
while "//" in collapsed:
collapsed = collapsed.replace("//", "/")
request.scope["path"] = collapsed
return await call_next(request)
app.add_middleware(CollapseDuplicatePathSlashesMiddleware)
app.add_middleware(
CORSMiddleware,
allow_origins=[
"https://nec.clients.septem.pro",
"https://localhost:3000",
"http://localhost:8501",
],
allow_credentials=True,
allow_methods=["*"],
allow_headers=["*"],
expose_headers=["Content-Disposition", "Content-Length"],
)
Instrumentator().instrument(app).expose(app)
# ------------------ Вспомогательные функции ------------------
def load_shipping_types():
if not os.path.exists(SHIPPING_TYPES_FILE):
return []
with open(SHIPPING_TYPES_FILE, "r", encoding="utf-8") as f:
types = json.load(f)
normalized_types = []
for t in types:
normalized = {}
for key, value in t.items():
clean_key = key.strip()
if clean_key == "keywords" and isinstance(value, list):
normalized[clean_key] = [kw.strip().strip('"').strip("'") for kw in value if kw.strip()]
elif isinstance(value, str):
normalized[clean_key] = value.strip()
else:
normalized[clean_key] = value
normalized_types.append(normalized)
return normalized_types
def save_shipping_types(types):
with open(SHIPPING_TYPES_FILE, "w", encoding="utf-8") as f:
json.dump(types, f, ensure_ascii=False, indent=2)
def load_processed_shipping_criteria() -> Dict[str, str]:
if not os.path.exists(PROCESSED_SHIPPING_TYPES_FILE):
return {}
try:
with open(PROCESSED_SHIPPING_TYPES_FILE, "r", encoding="utf-8") as f:
data = json.load(f)
if isinstance(data, dict):
return {str(k): str(v) for k, v in data.items() if v is not None}
return {}
except Exception as e:
logger.warning(f"Failed to load processed shipping criteria: {e}")
return {}
def save_processed_shipping_criteria(mapping: Dict[str, str]):
with open(PROCESSED_SHIPPING_TYPES_FILE, "w", encoding="utf-8") as f:
json.dump(mapping, f, ensure_ascii=False, indent=2)
def upsert_processed_shipping_criteria(type_name: str, processed_criteria: str):
if not isinstance(type_name, str) or not type_name.strip():
return
processed_criteria = processed_criteria if isinstance(processed_criteria, str) else ""
mapping = load_processed_shipping_criteria()
mapping[type_name] = processed_criteria
save_processed_shipping_criteria(mapping)
# ------------------ Эндпоинты ------------------
@app.get("/addin/taskpane.html")
async def get_taskpane():
if os.path.exists("addin/taskpane.html"):
return FileResponse("addin/taskpane.html")
elif os.path.exists("frontend/taskpane.html"):
return FileResponse("frontend/taskpane.html")
else:
return {"error": "taskpane.html not found"}
@app.get("/addin/commands.html")
async def get_commands():
if os.path.exists("addin/commands.html"):
return FileResponse("addin/commands.html")
elif os.path.exists("frontend/commands.html"):
return FileResponse("frontend/commands.html")
else:
return {"error": "commands.html not found"}
@app.post("/process-outlook-emails")
async def process_outlook_emails(request: OutlookEmailsRequest):
try:
# Генерируем session_id, если не передан
session_id = request.session_id or hashlib.md5(
f"{time.time()}{str(request.emails)}".encode()
).hexdigest()[:16]
# Сохраняем письма (со всеми полями) в хранилище
emails_list = [email.model_dump() for email in request.emails]
_session_emails[session_id] = emails_list
logger.info(f"Сессия {session_id}: сохранено {len(emails_list)} писем")
return {
"status": "success",
"session_id": session_id,
"emails_processed": len(request.emails),
}
except Exception as e:
logger.error(f"Process Outlook emails error: {e}")
raise HTTPException(status_code=500, detail=str(e))
@app.post("/query-cargo", response_model=CargoQueryResponse)
async def query_cargo(request: CargoQueryRequest):
try:
emails = _session_emails.get(request.session_id, [])
if not emails:
raise HTTPException(status_code=400, detail="No emails found for this session")
# Вызываем анализ
result = rag.analyze(emails, request.query)
if "error" in result:
raise HTTPException(status_code=500, detail=result["error"])
# Формируем текстовый ответ (answer) из критериев
answer_lines = []
for item in result.get("criteria_results", []):
val = item["value"]
if val is None:
val_str = "Не указано"
else:
val_str = str(val)
answer_lines.append(f"{item['criterion']}: {val_str}")
answer = "\n".join(answer_lines)
# Структурированные данные для UI
structured_data = {
"shipments": [
{
"shipping_type": result.get("shipping_type"),
"criteria_results": result.get("criteria_results"),
}
]
}
# Источники оригинальные письма (без вложений, только метаданные)
sources = []
for email in emails:
src = {
"id": email.get("id"),
"subject": email.get("subject"),
"sender": email.get("sender"),
"senderName": email.get("senderName"),
"receivedTime": email.get("receivedTime"),
"to": email.get("to"),
"cc": email.get("cc"),
}
sources.append(src)
return CargoQueryResponse(
answer=answer,
structured_data=structured_data,
sources=sources,
total_emails_analyzed=len(emails),
)
except HTTPException:
raise
except Exception as e:
logger.error(f"Cargo query error: {e}")
raise HTTPException(status_code=500, detail=str(e))
@app.post("/generate-cargo-report")
async def generate_cargo_report(session_id: str = Body(..., embed=True)):
try:
emails = _session_emails.get(session_id, [])
if not emails:
raise HTTPException(status_code=400, detail="No emails found for this session")
result = rag.analyze(emails)
if "error" in result:
raise HTTPException(status_code=500, detail=result["error"])
# Подготавливаем sources — все письма с вложениями (content_base64 обязательно!)
sources = []
for idx, email in enumerate(emails):
src = {
"id": email.get("id"),
"subject": email.get("subject", ""),
"sender": email.get("sender", ""),
"senderName": email.get("senderName", ""),
"receivedTime": email.get("receivedTime", ""),
"to": email.get("to", ""),
"cc": email.get("cc", ""),
"body": email.get("body", ""),
"body_html": email.get("body_html", ""),
"attachments": []
}
for att in email.get("attachments", []):
src["attachments"].append({
"filename": att.get("filename", ""),
"size": att.get("size", 0),
"text": att.get("text", ""),
"content_base64": att.get("content", "")
})
sources.append(src)
return {
"shipping_type": result.get("shipping_type", ""),
"criteria_results": result.get("criteria_results", []),
"emails_count": len(emails),
"sources": sources
}
except HTTPException:
raise
except Exception as e:
logger.error(f"Generate report error: {e}")
raise HTTPException(status_code=500, detail=str(e))
@app.get("/email-sessions/{session_id}")
async def get_session_info(session_id: str):
emails = _session_emails.get(session_id, [])
return {
"session_id": session_id,
"created_at": datetime.now().isoformat(),
"emails_count": len(emails),
}
@app.post("/upload", include_in_schema=False)
async def upload_document(file: UploadFile = File(...)):
try:
content = await file.read()
doc_id = hashlib.md5(file.filename.encode()).hexdigest()
# Здесь можно сохранить файл, но пока заглушка
return {"status": "processing", "document_id": doc_id}
except Exception as e:
logger.error(f"Upload error: {e}")
raise HTTPException(status_code=500, detail=str(e))
@app.get("/health")
async def health():
return {"status": "healthy"}
# ------------------ Вложения (старые методы, адаптированные под хранилище) ------------------
@app.get("/attachments/{session_id}/{email_index}/{attachment_index}")
async def get_attachment(session_id: str, email_index: int, attachment_index: int):
emails = _session_emails.get(session_id, [])
if email_index < 0 or email_index >= len(emails):
raise HTTPException(status_code=404, detail="Email not found")
attachments = emails[email_index].get("attachments", [])
if attachment_index < 0 or attachment_index >= len(attachments):
raise HTTPException(status_code=404, detail="Attachment not found")
att = attachments[attachment_index]
filename = att.get("filename", "attachment")
content_base64 = att.get("content")
def make_content_disposition(filename: str, disposition: str = "inline") -> str:
ascii_filename = filename.encode('ascii', 'ignore').decode('ascii') or 'attachment'
utf8_filename = quote(filename, safe='')
return f'{disposition}; filename="{ascii_filename}"; filename*=UTF-8\'\'{utf8_filename}'
if not content_base64:
# Если нет контента, возвращаем пустой ответ
return Response(
content="",
media_type="text/plain",
headers={
"Content-Disposition": make_content_disposition(filename + ".txt"),
"Access-Control-Expose-Headers": "Content-Disposition, Content-Length",
},
)
try:
file_content = base64.b64decode(content_base64)
ext = filename.lower().split('.')[-1] if '.' in filename else ''
mime_types = {
'pdf': 'application/pdf',
'doc': 'application/msword',
'docx': 'application/vnd.openxmlformats-officedocument.wordprocessingml.document',
'xls': 'application/vnd.ms-excel',
'xlsx': 'application/vnd.openxmlformats-officedocument.spreadsheetml.sheet',
'txt': 'text/plain',
'csv': 'text/csv',
'png': 'image/png',
'jpg': 'image/jpeg',
'jpeg': 'image/jpeg',
'gif': 'image/gif',
'bmp': 'image/bmp',
'zip': 'application/zip',
'rar': 'application/vnd.rar',
}
media_type = mime_types.get(ext, 'application/octet-stream')
return Response(
content=file_content,
media_type=media_type,
headers={
"Content-Disposition": make_content_disposition(filename),
"Content-Length": str(len(file_content)),
"Access-Control-Expose-Headers": "Content-Disposition, Content-Length",
},
)
except Exception as e:
logger.error(f"Error serving attachment: {e}")
raise HTTPException(status_code=500, detail=f"Error processing file: {str(e)}")
@app.get("/email-attachments/{session_id}")
async def get_email_attachments(session_id: str):
emails = _session_emails.get(session_id, [])
files = []
for email in emails:
for att in email.get("attachments", []):
filename = att.get("filename", "")
ext = filename.split(".")[-1].lower() if "." in filename else ""
if ext in ("png", "jpg", "jpeg", "gif", "bmp", "tiff", "webp"):
continue
if att.get("content"):
files.append({
"filename": filename,
"content_base64": att["content"],
})
return files
# ------------------ CRUD для типов перевозок ------------------
@app.get("/shipping-types", response_model=List[ShippingType])
async def get_shipping_types():
return load_shipping_types()
@app.post("/shipping-types", response_model=ShippingType)
async def create_shipping_type(item: ShippingTypeCreate):
types = load_shipping_types()
new_id = max([t["id"] for t in types], default=0) + 1
new_item = item.model_dump()
new_item["id"] = new_id
types.append(new_item)
save_shipping_types(types)
try:
type_name = new_item.get("name", "")
criteria_text = new_item.get("criteria", "") or ""
processed = rag.process_shipping_type_criteria(criteria_text)
upsert_processed_shipping_criteria(type_name, processed)
rag.reload_shipping_types()
except Exception as e:
logger.warning(f"AI criteria processing failed on create: {e}")
return new_item
@app.put("/shipping-types/{item_id}", response_model=ShippingType)
async def update_shipping_type(item_id: int, item: ShippingTypeUpdate):
types = load_shipping_types()
for t in types:
if t["id"] == item_id:
old_name = t.get("name", "")
t.update(item.model_dump())
save_shipping_types(types)
try:
type_name = t.get("name", "")
criteria_text = t.get("criteria", "") or ""
processed = rag.process_shipping_type_criteria(criteria_text)
if isinstance(old_name, str) and old_name.strip() and old_name != type_name:
mapping = load_processed_shipping_criteria()
if old_name in mapping:
mapping.pop(old_name, None)
save_processed_shipping_criteria(mapping)
upsert_processed_shipping_criteria(type_name, processed)
rag.reload_shipping_types()
except Exception as e:
logger.warning(f"AI criteria processing failed on update: {e}")
return t
raise HTTPException(status_code=404, detail="Type not found")
@app.delete("/shipping-types/{item_id}")
async def delete_shipping_type(item_id: int):
types = load_shipping_types()
removed = None
for t in types:
if t.get("id") == item_id:
removed = t
break
new_types = [t for t in types if t["id"] != item_id]
if len(new_types) == len(types):
raise HTTPException(status_code=404, detail="Type not found")
save_shipping_types(new_types)
try:
if removed:
type_name = removed.get("name", "")
mapping = load_processed_shipping_criteria()
if type_name in mapping:
mapping.pop(type_name, None)
save_processed_shipping_criteria(mapping)
rag.reload_shipping_types()
except Exception as e:
logger.warning(f"Failed to delete processed criteria on delete: {e}")
return {"ok": True}
@app.post("/record-cargo-learning")
async def record_cargo_learning(req: CargoLearningRequest):
try:
ok = rag.record_cargo_learning(
structured_data=req.structured_data,
session_id=req.session_id,
context_preview=req.context_preview,
notes=req.notes,
)
return {"status": "ok" if ok else "skipped", "stored": ok}
except Exception as e:
logger.error(f"record-cargo-learning error: {e}")
raise HTTPException(status_code=500, detail=str(e))
if __name__ == "__main__":
uvicorn.run(app, host="0.0.0.0", port=8000)