app/routers/search.py

from fastapi import APIRouter, Request, Form
from fastapi.responses import HTMLResponse, StreamingResponse
from fastapi.templating import Jinja2Templates
from langchain.vectorstores import FAISS
from langchain.embeddings import HuggingFaceEmbeddings
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
from app.modules.evaluation import Evaluation  # 평가 모듈 임포트
import os

router = APIRouter()

# 로컬 모델 로드 (HuggingFace T5 사용)
# tokenizer = AutoTokenizer.from_pretrained("gogamza/kobart-base-v2")
# model = AutoModelForSeq2SeqLM.from_pretrained("gogamza/kobart-base-v2")

# def summarize_text(text: str, max_length=150):
#     inputs = tokenizer.encode("summarize: " + text, return_tensors="pt", max_length=512, truncation=True)
#     outputs = model.generate(inputs, max_length=512, min_length=50, length_penalty=2.0, num_beams=4, early_stopping=True)
#     return tokenizer.decode(outputs[0], skip_special_tokens=True)

# 템플릿 설정
templates = Jinja2Templates(directory="app/templates")

# 업로드된 파일로부터 생성된 벡터 데이터베이스 경로 설정
UPLOAD_DIRECTORY = "./uploaded_files"
db_path = os.path.join(UPLOAD_DIRECTORY, "faiss_index")

# 벡터 데이터베이스 로드
embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2")

@router.get("/search", response_class=HTMLResponse)
async def search_page(request: Request):
    return templates.TemplateResponse("search.html", {"request": request})


@router.post("/search/stream")
async def search_stream(query: str = Form(...), model: str = Form(...)):
     # 벡터 데이터베이스 로드
    vector_db = FAISS.load_local(db_path, embeddings, allow_dangerous_deserialization=True)
    results = vector_db.similarity_search(query, k=5)
    
    # 검색된 결과 텍스트를 연결
    full_text = "\n\n".join([result.page_content for result in results])
    
    # 평가 모듈 초기화
    evaluator = Evaluation(model=model)
    
    # 평가 기준
    instruction = (
        f"다음 텍스트를 기반으로 {query}\n\n{full_text}"
    )
    

    async def stream():
        async for data in evaluator.evaluate_stream(instruction):
            yield data

    return StreamingResponse(stream(), media_type="text/plain")