File size: 1,151 Bytes
e280b6e
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
from motor.motor_asyncio import AsyncIOMotorClient
from app.config import settings

class MongoDB:
    client: AsyncIOMotorClient = None
    db = None

    def connect(self) -> None:
        """Establish MongoDB connection client."""
        self.client = AsyncIOMotorClient(settings.MONGODB_URI)
        self.db = self.client[settings.MONGODB_DB]

    def disconnect(self) -> None:
        """Close MongoDB connection client."""
        if self.client:
            self.client.close()

    @property
    def email_bodies(self):
        return self.db["email_bodies"]

    @property
    def agent_memory(self):
        return self.db["agent_memory"]

    @property
    def agent_logs(self):
        return self.db["agent_logs"]

# Global MongoDB wrapper
mongo_db = MongoDB()

async def init_mongo() -> None:
    """Initialize collections and indexes."""
    mongo_db.connect()
    # Create indexes for high performance querying
    await mongo_db.email_bodies.create_index("email_id", unique=True)
    await mongo_db.agent_memory.create_index([("user_id", 1), ("email_id", 1)])
    await mongo_db.agent_logs.create_index([("run_id", 1), ("step", 1)])