from keble_helpers import PydanticModelConfig
from motor.motor_asyncio import AsyncIOMotorClient
from neo4j import AsyncSession
from pydantic import BaseModel, SkipValidation
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()

    amongo: SkipValidation[AsyncIOMotorClient]
    aredis: SkipValidation[AsyncRedis]
    extended_aredis: SkipValidation[ExtendedAsyncRedis]
    aneo4j: SkipValidation[AsyncSession]
    aqdrant: SkipValidation[AsyncQdrantClient]
