phoenix-oss/llama-stack-mirror
0
0
Fork 1
mirror of https://github.com/meta-llama/llama-stack.git synced 2025-08-02 08:44:44 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

fixed port and collection_name issue

Browse source
This commit is contained in:
shrinitgoyal 2024-11-28 22:03:30 +05:30
parent 23df7db896
commit c2f3a7f87e
2 changed files with 5 additions and 14 deletions
Download patch file Download diff file Expand all files Collapse all files

4
llama_stack/providers/utils/kvstore/config.py
View file

@ -17,7 +17,7 @@ from llama_stack.distribution.utils.config_dirs import RUNTIME_BASE_DIR
class KVStoreType(Enum): class KVStoreType(Enum):
redis = "redis" redis = "redis"
sqlite = "sqlite" sqlite = "sqlite"
postgres = "postgres", postgres = "postgres"
mongodb = "mongodb" mongodb = "mongodb"
@ -110,7 +110,7 @@ class PostgresKVStoreConfig(CommonConfig):
class MongoDBKVStoreConfig(CommonConfig): class MongoDBKVStoreConfig(CommonConfig):
type: Literal[KVStoreType.mongodb.value] = KVStoreType.mongodb.value type: Literal[KVStoreType.mongodb.value] = KVStoreType.mongodb.value
host: str = "localhost" host: str = "localhost"
port: int = 5432 port: int = 27017
db: str = "llamastack" db: str = "llamastack"
user: str = None user: str = None
password: Optional[str] = None password: Optional[str] = None

15
llama_stack/providers/utils/kvstore/mongodb/mongodb.py
View file

@ -1,4 +1,3 @@
import datetime
import logging import logging
from datetime import datetime from datetime import datetime
from typing import Optional, List from typing import Optional, List
@ -25,11 +24,11 @@ class MongoDBKVStoreImpl(KVStore):
"username": self.config.user, "username": self.config.user,
"password": self.config.password, "password": self.config.password,
} }
conn_creds = {k: v for k, v in conn_creds if v is not None} conn_creds = {k: v for k, v in conn_creds.items() if v is not None}
try: try:
self.conn = MongoClient(**conn_creds) self.conn = MongoClient(**conn_creds)
self.collection = self.conn[self.config.db] self.collection = self.conn[self.config.db][self.config.collection_name]
except (ConnectionError, ConfigurationError) as e: except (ConnectionError, ConfigurationError) as e:
raise Exception(f"Failed to connect to MongoDB: {e}") raise Exception(f"Failed to connect to MongoDB: {e}")
except Exception as e: except Exception as e:
@ -60,11 +59,7 @@ class MongoDBKVStoreImpl(KVStore):
async def get(self, key: str) -> Optional[str]: async def get(self, key: str) -> Optional[str]:
key = self._namespaced_key(key) key = self._namespaced_key(key)
query = { query = {
"key": key, "key": key
"$or": [
{"expiration": {"$exists": False}},
{"expiration": {"$gt": datetime.now(datetime.UTC)}},
],
} }
result = self.collection.find_one(query, {"value": 1, "_id": 0}) result = self.collection.find_one(query, {"value": 1, "_id": 0})
return result["value"] if result else None return result["value"] if result else None
@ -78,10 +73,6 @@ class MongoDBKVStoreImpl(KVStore):
end_key = self._namespaced_key(end_key) end_key = self._namespaced_key(end_key)
query = { query = {
"key": {"$gte": start_key, "$lt": end_key}, "key": {"$gte": start_key, "$lt": end_key},
"$or": [
{"expiration": {"$exists": False}},
{"expiration": {"$gt": datetime.now(datetime.UTC)}},
],
} }
cursor = self.collection.find(query, {"value": 1, "_id": 0}).sort("key", 1) cursor = self.collection.find(query, {"value": 1, "_id": 0}).sort("key", 1)
return [doc["value"] for doc in cursor] return [doc["value"] for doc in cursor]