from keble_helpers import PydanticModelConfig
from motor.motor_asyncio import AsyncIOMotorClient
from neo4j import AsyncTransaction
from pydantic import BaseModel
from qdrant_client import AsyncQdrantClient
from redis.asyncio import Redis as AsyncRedis

from keble_db.wrapper import ExtendedAsyncRedis


class AgentDbDeps(BaseModel):
    model_config = PydanticModelConfig.default(arbitrary_types_allowed=True)

    amongo: AsyncIOMotorClient
    aredis: AsyncRedis
    extended_aredis: ExtendedAsyncRedis
    aneo4j: AsyncTransaction
    aqdrant: AsyncQdrantClient
