| |
| |
| |
|
|
| from fastapi import FastAPI, Header, HTTPException |
| from pydantic import BaseModel |
| from typing import List, Optional |
| from utils.s3_utils import read_s3_file |
| from utils.embedding_utils import read_document, cumulative_semantic_chunking, embed_chunks |
| from utils.qdrant_utils import store_embeddings |
| import logging |
| import time |
| import os |
| from dotenv import load_dotenv |
|
|
| |
|
|
| |
| API_KEY = os.getenv('X_API_KEY') |
| bucket_name = os.getenv('bucket_name') |
|
|
| |
| logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s') |
|
|
| app = FastAPI() |
|
|
| class Metadata(BaseModel): |
| mime_type: str |
| file_size_bytes: str |
| file_format: str |
|
|
| class DocumentIndexRequest(BaseModel): |
| metadata: Metadata |
| bucket_key: str |
| user_id: str |
| org_id: str |
| file_id: str |
| data_source_id: int |
| |
| @app.get("/") |
| async def root(): |
| return {"message": "Welcome to the Document Indexing API!"} |
|
|
| @app.post("/api/document-index") |
| async def document_index(request: DocumentIndexRequest, x_api_key: str = Header(...)): |
| logging.info(f"Received request: {request}") |
| start_time = time.time() |
| |
| if x_api_key != API_KEY: |
| logging.info(x_api_key) |
| logging.info(API_KEY) |
| logging.warning("Unauthorized access attempt with invalid API key.") |
| raise HTTPException(status_code=401, detail="Unauthorized") |
| |
| try: |
| |
| |
| |
| content, metadata, file_format = read_s3_file(bucket_name, request.bucket_key) |
| logging.info(f"File {request.bucket_key} retrieved from S3 with format {file_format}.") |
|
|
|
|
| |
| text_content = read_document(content, file_id=request.file_id, file_format=file_format) |
| logging.info(f"Text content extracted from file {request.bucket_key}.") |
| |
|
|
| |
| chunks = cumulative_semantic_chunking(text_content, max_chunk_size=2048, similarity_threshold=0.6) |
| logging.info(f"Text content chunked into {len(chunks)} chunks.") |
| |
| embeddings, total_tokens = embed_chunks(chunks) |
| logging.info(f"Text content embedded into vectors. Total tokens used: {total_tokens}.") |
|
|
| |
| store_embeddings( |
| chunks=chunks, |
| embeddings=embeddings, |
| user_id=request.user_id, |
| data_source_id=request.data_source_id, |
| file_id=request.file_id, |
| organization_id=request.org_id, |
| s3_bucket_key=request.bucket_key, |
| total_tokens=total_tokens |
| ) |
| |
| logging.info(f"Embeddings for {request.bucket_key} stored successfully with metadata: {metadata}") |
| logging.info(f"Embeddings for stored successfully with {total_tokens} tokens") |
| time_taken = time.time() - start_time |
| logging.info(f"Time taken to process and embed the document: {time_taken} seconds") |
|
|
|
|
| return {"message": "Embeddings stored successfully"} |
| |
| |
| except FileNotFoundError as e: |
| logging.error(f"File not found: {str(e)}") |
| raise HTTPException(status_code=404, detail=str(e)) |
| except PermissionError as e: |
| logging.error(f"Permission error: {str(e)}") |
| raise HTTPException(status_code=403, detail=str(e)) |
| except HTTPException as e: |
| logging.error(f"HTTP error: {str(e.detail)}") |
| raise |
| except Exception as e: |
| logging.error(f"Error processing file {request.bucket_key}: {str(e)}") |
| raise HTTPException(status_code=500, detail=str(e)) |
|
|
| |
| if __name__ == "__main__": |
| import uvicorn |
| uvicorn.run(app, host="0.0.0.0", port=8000) |
|
|