api/function.py

import json
import logging
from typing import Any

from dotenv import load_dotenv
from fastapi import HTTPException, UploadFile
from fastapi.responses import JSONResponse

from core.chat.engine import Engine
from core.parser import clean_text
from langfuse.llama_index import LlamaIndexCallbackHandler
from script.document_uploader import Uploader
from script.vector_db import IndexManager
from service.aws_loader import Loader
from service.dto import BotResponseStreaming


load_dotenv()

# Configure logging
logging.basicConfig(level=logging.INFO)


async def data_ingestion(reference, file: UploadFile) -> Any:
    try:
        # Assuming you have a Langfuse callback handler
        langfuse_callback_handler = LlamaIndexCallbackHandler()
        langfuse_callback_handler.set_trace_params(
            user_id="admin_book_uploaded",
        )
        
        uploader = Uploader(reference, file)
        nodes_with_metadata, file_stream = await uploader.process_documents()

        # Build indexes using IndexManager
        index = IndexManager()
        index.build_indexes(nodes_with_metadata)

        
        # # Upload to AWS
        file_name = f"{reference['title']}"
        aws_loader = Loader()

        # file_obj = file
        aws_loader.upload_to_s3(file_stream, file_name)


        return json.dumps(
            {"status": "success", "message": "Vector Index loaded successfully."}
        )

    except Exception as e:
        # Log the error and raise HTTPException for FastAPI
        logging.error("An error occurred in data ingestion: %s", e)
        return JSONResponse(
            status_code=500,
            content="An internal server error occurred in data ingestion.",
        )

async def generate_streaming_completion(user_request, session_id):
    try:
        engine = Engine()
        index_manager = IndexManager()

        # Load existing indexes
        index = index_manager.load_existing_indexes()

        # Retrieve the chat engine with the loaded index
        chat_engine = engine.get_chat_engine(index, session_id)
        # Generate completion response
        response = chat_engine.stream_chat(user_request)

        completed_response = ""

        for gen in response.response_gen:
            completed_response += gen  # Concatenate the new string
            yield BotResponseStreaming(
                content=gen, completed_content=completed_response
            )

        nodes = response.source_nodes
        for node in nodes:
            reference = str(clean_text(node.node.get_text()))
            metadata = dict(node.node.metadata)
            score = float(node.score)
            yield BotResponseStreaming(
                completed_content=completed_response,
                reference=reference,
                metadata=metadata,
                score=score,
            )
    except Exception as e:
        yield {"error": str(e)}

    except Exception as e:
        # Log the error and raise HTTPException for FastAPI
        logging.error(f"An error occurred in generate text: {e}")
        raise HTTPException(
            status_code=500,
            detail="An internal server error occurred in generate text.",
        ) from e