server/routes/database.py

import xml.etree.ElementTree as ET
from elasticsearch import Elasticsearch, helpers
import torch
import numpy as np


from fastapi import APIRouter, Depends, HTTPException, Body, UploadFile, File
from server.utils.database import get_db
from fastapi.encoders import jsonable_encoder
import json
from server.utils.insert_data import insert_data_from_xml
from server.utils.search_data import search_data_from_database

router = APIRouter(prefix="/patent_data")


@router.post("/insert")
async def insert_data_from_xml(
    xml_file: UploadFile = File(...),
    db: Elasticsearch = Depends(get_db),
    index_name: str = "patents"
):
    try:
        # Read the uploaded file as bytes and decode to string
        xml_content = await xml_file.read()
        import tempfile
        with tempfile.NamedTemporaryFile(delete=False, suffix='.xml') as tmp:
            tmp.write(xml_content)
            tmp_path = tmp.name
        result = insert_data_from_xml(tmp_path, db=db, index_name=index_name)
        return {"message": result}
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))


@router.get("/search_text/{query}")
async def search_text(
    query: str,
    top_k: int = 5,
    db: Elasticsearch = Depends(get_db),
    index_name: str = "patents"
):
    try:
        results = search_data_from_database(
            query=query,
            top_k=top_k,
            db=db,
            index_name=index_name
        )
        return {"results": results}
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))
    


@router.post("/search_image")
async def search_image(
    image: UploadFile = File(...),
    top_k: int = 5,
    db: Elasticsearch = Depends(get_db),
    index_name: str = "patents"
):
    try:
        # Save the uploaded image to a temp file
        import tempfile
        image_bytes = await image.read()
        with tempfile.NamedTemporaryFile(delete=False, suffix='.jpg') as tmp:
            tmp.write(image_bytes)
            tmp_path = tmp.name
        results = search_data_from_database(
            image_path=tmp_path,
            top_k=top_k,
            db=db,
            index_name=index_name
        )
        return {"results": results}
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))