Add KV cache configuration and conditional cache router registration (#280)

Author: Adityav369

core/api.py

@@ -307,8 +307,9 @@ def _extract_provider(model_name: str) -> str:
 # Register logs router
 app.include_router(logs_router)
 
-# Register cache router
-app.include_router(cache_router)
+# Register cache router (only if KV cache is enabled)
+if settings.KV_CACHE_ENABLED:
+    app.include_router(cache_router)
 
 # Register graph router
 app.include_router(graph_router)

core/config.py

@@ -118,6 +118,9 @@ class Settings(BaseSettings):
     CACHE_CHUNK_MAX_BYTES: int = 10 * 1024 * 1024 * 1024
     CACHE_PATH: str = "./storage/cache"
 
+    # KV Cache configuration (LlamaCache for semantic caching)
+    KV_CACHE_ENABLED: bool = False
+
     # Vector store configuration
     VECTOR_STORE_PROVIDER: Literal["pgvector"]
     VECTOR_STORE_DATABASE_NAME: Optional[str] = None
@@ -475,4 +478,8 @@ def get_settings() -> Settings:
                 )
             settings_dict["TURBOPUFFER_API_KEY"] = os.environ["TURBOPUFFER_API_KEY"]
 
+    # Load kv_cache config
+    if "kv_cache" in config:
+        settings_dict["KV_CACHE_ENABLED"] = config["kv_cache"].get("enabled", False)
+
     return Settings(**settings_dict)

core/database/postgres_database.py

@@ -258,6 +258,27 @@ def __init__(
             f"max_overflow={max_overflow}, pool_recycle={pool_recycle}s"
         )
 
+        # Strip parameters that asyncpg doesn't accept as keyword arguments
+        # These will raise "unexpected keyword argument" errors
+        from urllib.parse import parse_qs, urlencode, urlparse, urlunparse
+
+        parsed = urlparse(uri)
+        query_params = parse_qs(parsed.query)
+
+        # List of parameters that asyncpg doesn't accept
+        incompatible_params = ["sslmode", "channel_binding"]
+        removed_params = []
+
+        for param in incompatible_params:
+            if param in query_params:
+                query_params.pop(param, None)
+                removed_params.append(param)
+
+        if removed_params:
+            logger.debug(f"Removing parameters from PostgreSQL URI (not compatible with asyncpg): {removed_params}")
+            parsed = parsed._replace(query=urlencode(query_params, doseq=True))
+            uri = urlunparse(parsed)
+
         # Create async engine with explicit pool settings
         self.engine = create_async_engine(
             uri,

core/services_init.py

@@ -16,12 +16,10 @@
 from pathlib import Path
 from typing import Optional
 
-from core.cache.llama_cache_factory import LlamaCacheFactory
 from core.completion.litellm_completion import LiteLLMCompletionModel
 from core.config import get_settings
 from core.database.postgres_database import PostgresDatabase
 from core.embedding.colpali_api_embedding_model import ColpaliApiEmbeddingModel
-from core.embedding.colpali_embedding_model import ColpaliEmbeddingModel
 from core.embedding.litellm_embedding import LiteLLMEmbeddingModel
 from core.parser.morphik_parser import MorphikParser
 from core.reranker.flag_reranker import FlagReranker
@@ -118,7 +116,14 @@
 # Cache factory
 # ---------------------------------------------------------------------------
 
-cache_factory = LlamaCacheFactory(Path(settings.STORAGE_PATH))
+cache_factory = None
+if settings.KV_CACHE_ENABLED:
+    from core.cache.llama_cache_factory import LlamaCacheFactory
+
+    cache_factory = LlamaCacheFactory(Path(settings.STORAGE_PATH))
+    logger.info("KV cache enabled - initialized LlamaCacheFactory")
+else:
+    logger.info("KV cache disabled")
 
 # ---------------------------------------------------------------------------
 # ColPali multi-vector support
@@ -138,6 +143,8 @@
             colpali_vector_store = None
         case "local":
             logger.info("Initializing ColPali in local mode")
+            from core.embedding.colpali_embedding_model import ColpaliEmbeddingModel
+
             colpali_embedding_model = ColpaliEmbeddingModel()
             # Choose multivector store implementation based on provider and dual ingestion setting
             if settings.ENABLE_DUAL_MULTIVECTOR_INGESTION:

morphik.docker.toml

@@ -128,6 +128,9 @@ provider = "pgvector"
 [multivector_store]
 provider = "postgres"
 
+[kv_cache]
+enabled = false  # Enable LlamaCache for semantic caching
+
 [redis]
 url = "redis://redis:6379/0"  # Docker service name
 host = "redis"

morphik.toml

@@ -134,6 +134,9 @@ provider = "pgvector"
 [multivector_store]
 provider = "postgres"  # "morphik" # "postgres"  # "morphik" # "postgres"  # "postgres" or "morphik" for fast implementation
 
+[kv_cache]
+enabled = false  # Enable LlamaCache for semantic caching
+
 [redis]
 url = "redis://localhost:6379/0"  # Full Redis URL (takes precedence over host/port)
 host = "localhost"