From c2f3a7f87e186f5e5d7612f40a111d47b8bb34a7 Mon Sep 17 00:00:00 2001
From: shrinitgoyal <shrinit.goyal@engati.com>
Date: Thu, 28 Nov 2024 22:03:30 +0530
Subject: [PATCH] fixed port and collection_name issue

---
 llama_stack/providers/utils/kvstore/config.py     |  4 ++--
 .../providers/utils/kvstore/mongodb/mongodb.py    | 15 +++------------
 2 files changed, 5 insertions(+), 14 deletions(-)

diff --git a/llama_stack/providers/utils/kvstore/config.py b/llama_stack/providers/utils/kvstore/config.py
index 867ae3f98..b54926d07 100644
--- a/llama_stack/providers/utils/kvstore/config.py
+++ b/llama_stack/providers/utils/kvstore/config.py
@@ -17,7 +17,7 @@ from llama_stack.distribution.utils.config_dirs import RUNTIME_BASE_DIR
 class KVStoreType(Enum):
     redis = "redis"
     sqlite = "sqlite"
-    postgres = "postgres",
+    postgres = "postgres"
     mongodb = "mongodb"
 
 
@@ -110,7 +110,7 @@ class PostgresKVStoreConfig(CommonConfig):
 class MongoDBKVStoreConfig(CommonConfig):
     type: Literal[KVStoreType.mongodb.value] = KVStoreType.mongodb.value
     host: str = "localhost"
-    port: int = 5432
+    port: int = 27017
     db: str = "llamastack"
     user: str = None
     password: Optional[str] = None
diff --git a/llama_stack/providers/utils/kvstore/mongodb/mongodb.py b/llama_stack/providers/utils/kvstore/mongodb/mongodb.py
index f749f26d0..2aaf3ec9f 100644
--- a/llama_stack/providers/utils/kvstore/mongodb/mongodb.py
+++ b/llama_stack/providers/utils/kvstore/mongodb/mongodb.py
@@ -1,4 +1,3 @@
-import datetime
 import logging
 from datetime import datetime
 from typing import Optional, List
@@ -25,11 +24,11 @@ class MongoDBKVStoreImpl(KVStore):
                 "username": self.config.user,
                 "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:
                 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:
                 raise Exception(f"Failed to connect to MongoDB: {e}")
         except Exception as e:
@@ -60,11 +59,7 @@ class MongoDBKVStoreImpl(KVStore):
     async def get(self, key: str) -> Optional[str]:
         key = self._namespaced_key(key)
         query = {
-            "key": key,
-            "$or": [
-                {"expiration": {"$exists": False}},
-                {"expiration": {"$gt": datetime.now(datetime.UTC)}},
-            ],
+            "key": key
         }
         result = self.collection.find_one(query, {"value": 1, "_id": 0})
         return result["value"] if result else None
@@ -78,10 +73,6 @@ class MongoDBKVStoreImpl(KVStore):
         end_key = self._namespaced_key(end_key)
         query = {
             "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)
         return [doc["value"] for doc in cursor]