chat / app /app.py
ariansyahdedy's picture
Add prompt edit and api key config
8d2f9d4
from fastapi import FastAPI, Request, status
from fastapi.responses import JSONResponse
from fastapi.responses import Response
from fastapi.exceptions import HTTPException
from slowapi import Limiter, _rate_limit_exceeded_handler
from slowapi.errors import RateLimitExceeded
from slowapi.util import get_remote_address
from slowapi.middleware import SlowAPIMiddleware
from typing import Dict, List
from prometheus_client import Counter, Histogram, start_http_server
from pydantic import BaseModel, ValidationError
from app.services.message import generate_reply, send_reply
import logging
from datetime import datetime
from sentence_transformers import SentenceTransformer
from contextlib import asynccontextmanager
# from app.db.database import create_indexes, init_db
from app.services.webhook_handler import verify_webhook
from app.handlers.message_handler import MessageHandler
from app.handlers.webhook_handler import WebhookHandler
from app.handlers.media_handler import WhatsAppMediaHandler
from app.services.cache import MessageCache
from app.services.chat_manager import ChatManager
from app.api.api_file import file_router
from app.utils.load_env import ACCESS_TOKEN
from app.search.rag_pipeline import RAGSystem
# Configure logging
logging.basicConfig(
level=logging.INFO,
format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger(__name__)
# Initialize handlers at startup
message_handler = None
webhook_handler = None
async def setup_message_handler():
logger = logging.getLogger(__name__)
message_cache = MessageCache()
chat_manager = ChatManager()
media_handler = WhatsAppMediaHandler()
return MessageHandler(
message_cache=message_cache,
chat_manager=chat_manager,
media_handler=media_handler,
logger=logger
)
async def setup_rag_system():
embedding_model = SentenceTransformer('all-MiniLM-L6-v2') # Replace with your model if different
rag_system = RAGSystem(embedding_model)
return rag_system
# Initialize FastAPI app
@asynccontextmanager
async def lifespan(app: FastAPI):
try:
# await init_db()
logger.info("Connected to the MongoDB database!")
rag_system = await setup_rag_system()
app.state.rag_system = rag_system
global message_handler, webhook_handler
message_handler = await setup_message_handler()
webhook_handler = WebhookHandler(message_handler)
# collections = app.database.list_collection_names()
# print(f"Collections in {db_name}: {collections}")
yield
except Exception as e:
logger.error(e)
# Initialize Limiter and Prometheus Metrics
limiter = Limiter(key_func=get_remote_address)
app = FastAPI(lifespan=lifespan)
app.state.limiter = limiter
app.add_exception_handler(RateLimitExceeded, _rate_limit_exceeded_handler)
# Add SlowAPI Middleware
app.add_middleware(SlowAPIMiddleware)
# app.include_router(users.router, prefix="/users", tags=["Users"])
app.include_router(file_router, prefix="/file_load", tags=["File Load"])
# Prometheus metrics
webhook_requests = Counter('webhook_requests_total', 'Total webhook requests')
webhook_processing_time = Histogram('webhook_processing_seconds', 'Time spent processing webhook')
# Pseudocode: You may have a database or environment variable to validate API keys
VALID_API_KEYS = {"!@#$%^"}
class WebhookPayload(BaseModel):
entry: List[Dict]
@app.post("/api/v1/messages")
@limiter.limit("20/minute")
async def process_message(request: Request):
try:
# Validate developer’s API key
api_key = request.headers.get("Authorization")
if not api_key or not api_key.startswith("Bearer "):
raise HTTPException(status_code=401, detail="Missing or invalid API key")
api_key_value = api_key.replace("Bearer ", "")
if api_key_value not in VALID_API_KEYS:
raise HTTPException(status_code=403, detail="Forbidden")
payload = await request.json()
# Extract needed credentials from query params or request body
# e.g., whatsapp_token, verify_token, llm_api_key, llm_model
whatsapp_token = request.query_params.get("whatsapp_token")
whatsapp_url = request.query_params.get("whatsapp_url")
gemini_api = request.query_params.get("gemini_api")
llm_model = request.query_params.get("cx_code")
print(f"payload: {payload}")
response = await webhook_handler.process_webhook(
payload=payload,
whatsapp_token=whatsapp_token,
whatsapp_url=whatsapp_url,
gemini_api=gemini_api,
)
return JSONResponse(
content=response.__dict__,
status_code=status.HTTP_200_OK
)
except ValidationError as ve:
logger.error(f"Validation error: {ve}")
return JSONResponse(
content={"status": "error", "detail": ve.errors()},
status_code=status.HTTP_422_UNPROCESSABLE_ENTITY
)
except Exception as e:
logger.error(f"Unexpected error: {str(e)}")
return JSONResponse(
content={"status": "error", "detail": str(e)},
status_code=status.HTTP_500_INTERNAL_SERVER_ERROR
)
app.get("/webhook")(verify_webhook)