added upstash

pinecone_keep_alive_service.py

@@ -3,6 +3,8 @@ import logging
 from datetime import datetime
 from dotenv import load_dotenv
 from pinecone import Pinecone
+from pydantic import BaseModel
+from typing import Any
 
 # Configure logging
 logging.basicConfig(
@@ -17,8 +19,23 @@ load_dotenv()
 keys = os.getenv("PINECONE_API_KEYS", "").split(",")
 
 
-def ping_all_pinecone_indexes():
-    results = []
+
+class PineconePingResponse(BaseModel):
+    service_name: str
+    success: bool
+    error: str | None = None
+    time: str
+    api_key: str
+    index: str | None = None
+    stats: dict[str, Any] | None = None
+
+
+class PineconePingAllResponse(BaseModel):
+    services: list[PineconePingResponse]
+
+
+def ping_all_pinecone_indexes() -> PineconePingAllResponse:
+    results: list[PineconePingResponse] = []
 
     for key in filter(None, (k.strip() for k in keys)):
         pc = Pinecone(api_key=key)
@@ -28,39 +45,47 @@ def ping_all_pinecone_indexes():
             indexes = pc.list_indexes()
         except Exception as e:
             logger.error(f"[{short_key}] failed to list indexes: {e}")
-            results.append({
-                "service_name": "pinecone",
-                "success": False,
-                "error": str(e),
-                "time": datetime.utcnow().isoformat(),
-                "api_key": short_key,
-                "index": None,
-            })
+            results.append(
+                PineconePingResponse(
+                    service_name="pinecone",
+                    success=False,
+                    error=str(e),
+                    time=datetime.utcnow().isoformat(),
+                    api_key=short_key,
+                )
+            )
             continue
 
         for index in indexes:
             try:
                 stats = pc.Index(index.name).describe_index_stats()
-                logger.info(f"[{short_key}] ping {index.name} ✓")
 
-                results.append({
-                    "service_name": f"pinecone-{index.name}",
-                    "success": True,
-                    "error": None,
-                    "time": datetime.utcnow().isoformat(),
-                    "api_key": short_key,
-                    "index": index.name,
-                    "stats": stats,
-                })
+                # 🔑 Make sure it's plain dict
+                stats = dict(stats) if not isinstance(stats, dict) else stats
+
+                logger.info(f"[{short_key}] ping {index.name} ✓")
+                results.append(
+                    PineconePingResponse(
+                        service_name=f"pinecone-{index.name}",
+                        success=True,
+                        error=None,
+                        time=datetime.utcnow().isoformat(),
+                        api_key=short_key,
+                        index=index.name,
+                        stats=stats,
+                    )
+                )
             except Exception as e:
                 logger.error(f"[{short_key}] ping {index.name} ✗: {e}")
-                results.append({
-                    "service_name": f"pinecone-{index.name}",
-                    "success": False,
-                    "error": str(e),
-                    "time": datetime.utcnow().isoformat(),
-                    "api_key": short_key,
-                    "index": index.name,
-                })
-
-    return results
+                results.append(
+                    PineconePingResponse(
+                        service_name=f"pinecone-{index.name}",
+                        success=False,
+                        error=str(e),
+                        time=datetime.utcnow().isoformat(),
+                        api_key=short_key,
+                        index=index.name,
+                    )
+                )
+
+    return PineconePingAllResponse(services=results)

redis_keep_alive_service.py

@@ -10,8 +10,8 @@ from pydantic import BaseModel
 # Logging (console only)
 logging.basicConfig(
     level=logging.INFO,
-    format='%(asctime)s - %(levelname)s - %(message)s',
-    handlers=[logging.StreamHandler()]
+    format="%(asctime)s - %(levelname)s - %(message)s",
+    handlers=[logging.StreamHandler()],
 )
 logger = logging.getLogger(__name__)
 
@@ -29,7 +29,6 @@ REDIS_SERVICES = {
         "port": os.getenv("REDIS_DEVELOPMENT_PORT"),
         "password": os.getenv("REDIS_DEVELOPMENT_PASSWORD"),
     },
-
     # Upstash Redis (REST API)
     "upstash_dev": {
         "rest_url": os.getenv("UPSTASH_REDIS_DEV_REST_URL"),
@@ -42,7 +41,7 @@ REDIS_SERVICES = {
 }
 
 
-# ---------- Pydantic Response ----------
+# ---------- Pydantic Response Models ----------
 class RedisPingResponse(BaseModel):
     service_name: str
     success: bool
@@ -50,6 +49,13 @@ class RedisPingResponse(BaseModel):
     time: str
     latency_ms: float | None = None
     type: str  # "self-managed" or "upstash"
+    set_latency_ms: float | None = None
+    get_latency_ms: float | None = None
+    value: str | None = None
+
+
+class RedisPingAllResponse(BaseModel):
+    services: list[RedisPingResponse]
 
 
 # ---------- Self-managed Redis ----------
@@ -72,15 +78,17 @@ def ping_redis(service_name: str, redis_config: dict) -> RedisPingResponse:
         key = f"healthcheck:{uuid.uuid4()}"
         r.set(key, "ok")
         got = r.get(key)
+        got_value = got.decode() if got else None  # ✅ ensure only str, not bytes or Redis object
 
-        logger.info(f"{service_name} PING OK, latency {latency:.2f}ms, value={got.decode()}")
+        logger.info(f"{service_name} PING OK, latency {latency:.2f}ms, value={got_value}")
         return RedisPingResponse(
             service_name=service_name,
             success=True,
             error=None,
             time=now,
             latency_ms=latency,
-            type="self-managed"
+            type="self-managed",
+            value=got_value,
         )
     except Exception as e:
         error_msg = str(e)
@@ -91,13 +99,13 @@ def ping_redis(service_name: str, redis_config: dict) -> RedisPingResponse:
             error=error_msg,
             time=now,
             latency_ms=None,
-            type="self-managed"
+            type="self-managed",
         )
 
 
 # ---------- Upstash Redis ----------
-def ping_upstash(redis_config):
-    """Ping and benchmark Upstash Redis REST API."""
+def ping_upstash(service_name: str, redis_config: dict) -> RedisPingResponse:
+    now = datetime.utcnow().isoformat()
     try:
         headers = {"Authorization": f"Bearer {redis_config['rest_token']}"}
 
@@ -107,38 +115,32 @@ def ping_upstash(redis_config):
         resp = requests.get(url, headers=headers)
         latency = (datetime.now() - start_time).total_seconds() * 1000
 
-        if resp.status_code == 200:
-            try:
-                result = resp.json().get("result")
-            except Exception:
-                result = resp.text.strip()
-
-            if result and result.upper() == "PONG":
-                logger.info(
-                    f"PING {redis_config['rest_url']} "
-                    f"- Response: {result}, Latency: {latency:.2f}ms"
-                )
-            else:
-                logger.error(f"Unexpected Upstash PING response: {resp.status_code} - {resp.text}")
-                return {
-                    "service_name": redis_config["rest_url"],
-                    "success": False,
-                    "error": f"Unexpected response: {resp.status_code} - {resp.text}",
-                    "time": datetime.now().isoformat(),
-                    "latency_ms": None,
-                    "type": "upstash",
-                }
-        else:
-            return {
-                "service_name": redis_config["rest_url"],
-                "success": False,
-                "error": f"HTTP {resp.status_code}: {resp.text}",
-                "time": datetime.now().isoformat(),
-                "latency_ms": None,
-                "type": "upstash",
-            }
-
-        # ✅ If ping is OK, also test SET/GET
+        if resp.status_code != 200:
+            return RedisPingResponse(
+                service_name=service_name,
+                success=False,
+                error=f"HTTP {resp.status_code}: {resp.text}",
+                time=now,
+                latency_ms=None,
+                type="upstash",
+            )
+
+        try:
+            result = resp.json().get("result")
+        except Exception:
+            result = resp.text.strip()
+
+        if not result or result.upper() != "PONG":
+            return RedisPingResponse(
+                service_name=service_name,
+                success=False,
+                error=f"Unexpected response: {resp.status_code} - {resp.text}",
+                time=now,
+                latency_ms=None,
+                type="upstash",
+            )
+
+        # ✅ SET/GET test
         key = f"healthcheck:{uuid.uuid4()}"
         value = "ok"
 
@@ -154,44 +156,40 @@ def ping_upstash(redis_config):
         if get_resp.status_code == 200:
             got_value = get_resp.json().get("result")
 
-        return {
-            "service_name": redis_config.get("name", "upstash"),
-            "success": True,
-            "error": None,
-            "time": datetime.now().isoformat(),
-            "latency_ms": latency,
-            "type": "upstash",
-            "set_latency_ms": set_latency,
-            "get_latency_ms": get_latency,
-            "value": got_value,
-        }
-
+        return RedisPingResponse(
+            service_name=service_name,
+            success=True,
+            error=None,
+            time=now,
+            latency_ms=latency,
+            type="upstash",
+            set_latency_ms=set_latency,
+            get_latency_ms=get_latency,
+            value=got_value,
+        )
     except Exception as e:
-        logger.error(f"Failed to test Upstash: {str(e)}")
-        return {
-            "service_name": redis_config.get("name", "upstash"),
-            "success": False,
-            "error": str(e),
-            "time": datetime.now().isoformat(),
-            "latency_ms": None,
-            "type": "upstash",
-        }
+        return RedisPingResponse(
+            service_name=service_name,
+            success=False,
+            error=str(e),
+            time=now,
+            latency_ms=None,
+            type="upstash",
+        )
 
 
 # ---------- Run All ----------
-def ping_all_redis_projects() -> list[RedisPingResponse]:
+def ping_all_redis_projects() -> RedisPingAllResponse:
     logger.info("Starting Redis service health check")
     results: list[RedisPingResponse] = []
 
     for service_name, config in REDIS_SERVICES.items():
         if not config or not any(config.values()):
-            continue  # skip empty configs
+            continue
         if "rest_url" in config:  # Upstash
-            results.append(ping_upstash(config))
+            results.append(ping_upstash(service_name, config))
         else:  # Self-managed
             results.append(ping_redis(service_name, config))
 
     logger.info("Completed all Redis service checks")
-    return results
-
-
+    return RedisPingAllResponse(services=results)