ariansyahdedy commited on
Commit
b0f8541
·
1 Parent(s): be9a762
Files changed (1) hide show
  1. app/main.py +11 -2
app/main.py CHANGED
@@ -12,7 +12,8 @@ from pydantic import BaseModel, ValidationError
12
  from app.services.message import generate_reply, send_reply
13
  import logging
14
  from datetime import datetime
15
- import time
 
16
  from contextlib import asynccontextmanager
17
  # from app.db.database import create_indexes, init_db
18
  from app.services.webhook_handler import verify_webhook
@@ -22,6 +23,7 @@ from app.handlers.media_handler import WhatsAppMediaHandler
22
  from app.services.cache import MessageCache
23
  from app.services.chat_manager import ChatManager
24
  from app.api.api_prompt import prompt_router
 
25
  from app.utils.load_env import ACCESS_TOKEN, WHATSAPP_API_URL, GEMINI_API
26
 
27
  # Configure logging
@@ -49,6 +51,11 @@ async def setup_message_handler():
49
  media_handler=media_handler,
50
  logger=logger
51
  )
 
 
 
 
 
52
  # Initialize FastAPI app
53
  @asynccontextmanager
54
  async def lifespan(app: FastAPI):
@@ -58,7 +65,9 @@ async def lifespan(app: FastAPI):
58
  # await init_db()
59
 
60
  logger.info("Connected to the MongoDB database!")
 
61
 
 
62
 
63
  global message_handler, webhook_handler
64
  message_handler = await setup_message_handler()
@@ -80,7 +89,7 @@ app.add_middleware(SlowAPIMiddleware)
80
 
81
  # app.include_router(users.router, prefix="/users", tags=["Users"])
82
  app.include_router(prompt_router, prefix="/prompts", tags=["Prompts"])
83
-
84
  # Prometheus metrics
85
  webhook_requests = Counter('webhook_requests_total', 'Total webhook requests')
86
  webhook_processing_time = Histogram('webhook_processing_seconds', 'Time spent processing webhook')
 
12
  from app.services.message import generate_reply, send_reply
13
  import logging
14
  from datetime import datetime
15
+ from sentence_transformers import SentenceTransformer
16
+ from app.search.rag_pipeline import RAGSystem
17
  from contextlib import asynccontextmanager
18
  # from app.db.database import create_indexes, init_db
19
  from app.services.webhook_handler import verify_webhook
 
23
  from app.services.cache import MessageCache
24
  from app.services.chat_manager import ChatManager
25
  from app.api.api_prompt import prompt_router
26
+ from app.api.api_file import file_router
27
  from app.utils.load_env import ACCESS_TOKEN, WHATSAPP_API_URL, GEMINI_API
28
 
29
  # Configure logging
 
51
  media_handler=media_handler,
52
  logger=logger
53
  )
54
+ async def setup_rag_system():
55
+ embedding_model = SentenceTransformer('all-MiniLM-L6-v2') # Replace with your model if different
56
+ rag_system = RAGSystem(embedding_model)
57
+
58
+ return rag_system
59
  # Initialize FastAPI app
60
  @asynccontextmanager
61
  async def lifespan(app: FastAPI):
 
65
  # await init_db()
66
 
67
  logger.info("Connected to the MongoDB database!")
68
+ rag_system = await setup_rag_system()
69
 
70
+ app.state.rag_system = rag_system
71
 
72
  global message_handler, webhook_handler
73
  message_handler = await setup_message_handler()
 
89
 
90
  # app.include_router(users.router, prefix="/users", tags=["Users"])
91
  app.include_router(prompt_router, prefix="/prompts", tags=["Prompts"])
92
+ app.include_router(file_router, prefix="/file_load", tags=["File Load"])
93
  # Prometheus metrics
94
  webhook_requests = Counter('webhook_requests_total', 'Total webhook requests')
95
  webhook_processing_time = Histogram('webhook_processing_seconds', 'Time spent processing webhook')