Added support for mongoDB KV store (#543)

Added the support for mongoDB as KV store
validated in mongodb, it is able to store agent data, session data and
turn data
<img width="1332" alt="image"
src="https://github.com/user-attachments/assets/867700a4-b9ee-4a3c-8278-f39074d39d56">
this is how run.yaml would look:
```
    config:
      persistence_store:
        type: mongodb
        namespace: null
        host: localhost
        port: 27017
        db: llamastack
        user: ""
        password: ""
        collection_name: llamastack_kvstore
```

---------

Co-authored-by: shrinitgoyal <shrinit.goyal@engati.com>
This commit is contained in:
Shrinit Goyal 2025-02-20 12:00:50 +05:30 committed by GitHub
parent 5966079770
commit b74f25035c
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
4 changed files with 106 additions and 2 deletions

View file

@ -18,6 +18,7 @@ class KVStoreType(Enum):
redis = "redis"
sqlite = "sqlite"
postgres = "postgres"
mongodb = "mongodb"
class CommonConfig(BaseModel):
@ -101,7 +102,30 @@ class PostgresKVStoreConfig(CommonConfig):
return v
class MongoDBKVStoreConfig(CommonConfig):
type: Literal[KVStoreType.mongodb.value] = KVStoreType.mongodb.value
host: str = "localhost"
port: int = 27017
db: str = "llamastack"
user: str = None
password: Optional[str] = None
collection_name: str = "llamastack_kvstore"
@classmethod
def sample_run_config(cls, collection_name: str = "llamastack_kvstore"):
return {
"type": "mongodb",
"namespace": None,
"host": "${env.MONGODB_HOST:localhost}",
"port": "${env.MONGODB_PORT:5432}",
"db": "${env.MONGODB_DB}",
"user": "${env.MONGODB_USER}",
"password": "${env.MONGODB_PASSWORD}",
"collection_name": "${env.MONGODB_COLLECTION_NAME:" + collection_name + "}",
}
KVStoreConfig = Annotated[
Union[RedisKVStoreConfig, SqliteKVStoreConfig, PostgresKVStoreConfig],
Union[RedisKVStoreConfig, SqliteKVStoreConfig, PostgresKVStoreConfig, MongoDBKVStoreConfig],
Field(discriminator="type", default=KVStoreType.sqlite.value),
]