database.py

import os
from datetime import datetime
from motor.motor_asyncio import AsyncIOMotorClient, AsyncIOMotorDatabase
from bson import ObjectId

client: AsyncIOMotorClient | None = None
db: AsyncIOMotorDatabase = None  # type: ignore


async def connect_db():
    global client, db
    uri = os.getenv("MONGODB_URI", "mongodb://admin:password123@mongo:27017/")
    db_name = os.getenv("MONGODB_DB", "research_agent")
    client = AsyncIOMotorClient(uri)
    db = client[db_name]

    # Create indexes for faster queries
    await db.research_history.create_index("created_at")
    await db.research_history.create_index("topic")
    print(f"Connected to MongoDB: {db_name}")


async def close_db():
    global client
    if client:
        client.close()
        print("MongoDB disconnected")


async def save_research(topic: str, depth: str, result: dict) -> str:
    doc = {
        "topic": topic,
        "depth": depth,
        "summary": result.get("summary", ""),
        "key_points": result.get("key_points", []),
        "sources": result.get("sources", []),
        "follow_up_questions": result.get("follow_up_questions", []),
        "warning": result.get("warning"),
        "created_at": datetime.utcnow(),
    }
    res = await db.research_history.insert_one(doc)
    return str(res.inserted_id)


async def get_history(limit: int = 20) -> list:
    cursor = db.research_history.find(
        {},
        {
            "_id": 1,
            "topic": 1,
            "depth": 1,
            "summary": 1,
            "created_at": 1,
        }
    ).sort("created_at", -1).limit(limit)

    results = []
    async for doc in cursor:
        doc["id"] = str(doc.pop("_id"))
        doc["created_at"] = doc["created_at"].isoformat()
        doc["summary"] = doc["summary"][:150] + "..." if len(doc.get("summary", "")) > 150 else doc.get("summary", "")
        results.append(doc)
    return results


async def get_by_id(research_id: str) -> dict | None:
    try:
        doc = await db.research_history.find_one({"_id": ObjectId(research_id)})
        if not doc:
            return None
        doc["id"] = str(doc.pop("_id"))
        doc["created_at"] = doc["created_at"].isoformat()
        return doc
    except Exception:
        return None


async def search_history(query: str) -> list:
    cursor = db.research_history.find(
        {"topic": {"$regex": query, "$options": "i"}},
        {
            "_id": 1,
            "topic": 1,
            "depth": 1,
            "summary": 1,
            "created_at": 1,
        }
    ).sort("created_at", -1).limit(10)

    results = []
    async for doc in cursor:
        doc["id"] = str(doc.pop("_id"))
        doc["created_at"] = doc["created_at"].isoformat()
        results.append(doc)
    return results


async def delete_by_id(research_id: str) -> bool:
    try:
        res = await db.research_history.delete_one({"_id": ObjectId(research_id)})
        return res.deleted_count == 1
    except Exception:
        return False