import json
import logging
import os
from abc import abstractmethod
from dataclasses import dataclass
from datetime import datetime, timedelta
from pathlib import Path
from typing import Any, Generic, Optional, TypeVar
import yaml
from core.base import Handler, generate_default_prompt_id
from .base import PostgresConnectionManager
logger = logging.getLogger(__name__)
T = TypeVar("T")
@dataclass
class CacheEntry(Generic[T]):
"""Represents a cached item with metadata."""
value: T
created_at: datetime
last_accessed: datetime
access_count: int = 0
class Cache(Generic[T]):
"""A generic cache implementation with TTL and LRU-like features."""
def __init__(
self,
ttl: Optional[timedelta] = None,
max_size: Optional[int] = 1000,
cleanup_interval: timedelta = timedelta(hours=1),
):
self._cache: dict[str, CacheEntry[T]] = {}
self._ttl = ttl
self._max_size = max_size
self._cleanup_interval = cleanup_interval
self._last_cleanup = datetime.now()
def get(self, key: str) -> Optional[T]:
"""Retrieve an item from cache."""
self._maybe_cleanup()
if key not in self._cache:
return None
entry = self._cache[key]
if self._ttl and datetime.now() - entry.created_at > self._ttl:
del self._cache[key]
return None
entry.last_accessed = datetime.now()
entry.access_count += 1
return entry.value
def set(self, key: str, value: T) -> None:
"""Store an item in cache."""
self._maybe_cleanup()
now = datetime.now()
self._cache[key] = CacheEntry(
value=value, created_at=now, last_accessed=now
)
if self._max_size and len(self._cache) > self._max_size:
self._evict_lru()
def invalidate(self, key: str) -> None:
"""Remove an item from cache."""
self._cache.pop(key, None)
def clear(self) -> None:
"""Clear all cached items."""
self._cache.clear()
def _maybe_cleanup(self) -> None:
"""Periodically clean up expired entries."""
now = datetime.now()
if now - self._last_cleanup > self._cleanup_interval:
self._cleanup()
self._last_cleanup = now
def _cleanup(self) -> None:
"""Remove expired entries."""
if not self._ttl:
return
now = datetime.now()
expired = [
k for k, v in self._cache.items() if now - v.created_at > self._ttl
]
for k in expired:
del self._cache[k]
def _evict_lru(self) -> None:
"""Remove least recently used item."""
if not self._cache:
return
lru_key = min(
self._cache.keys(), key=lambda k: self._cache[k].last_accessed
)
del self._cache[lru_key]
class CacheablePromptHandler(Handler):
"""Abstract base class that adds caching capabilities to prompt
handlers."""
def __init__(
self,
cache_ttl: Optional[timedelta] = timedelta(hours=1),
max_cache_size: Optional[int] = 1000,
):
self._prompt_cache = Cache[str](ttl=cache_ttl, max_size=max_cache_size)
self._template_cache = Cache[dict](
ttl=cache_ttl, max_size=max_cache_size
)
def _cache_key(
self, prompt_name: str, inputs: Optional[dict] = None
) -> str:
"""Generate a cache key for a prompt request."""
if inputs:
# Sort dict items for consistent keys
sorted_inputs = sorted(inputs.items())
return f"{prompt_name}:{sorted_inputs}"
return prompt_name
async def get_cached_prompt(
self,
prompt_name: str,
inputs: Optional[dict[str, Any]] = None,
prompt_override: Optional[str] = None,
bypass_cache: bool = False,
) -> str:
if prompt_override:
# If the user gave us a direct override, use it.
if inputs:
try:
return prompt_override.format(**inputs)
except KeyError:
return prompt_override
return prompt_override
cache_key = self._cache_key(prompt_name, inputs)
# If not bypassing, try returning from the prompt-level cache
if not bypass_cache:
cached = self._prompt_cache.get(cache_key)
if cached is not None:
logger.debug(f"Prompt cache hit: {cache_key}")
return cached
logger.debug(
"Prompt cache miss or bypass. Retrieving from DB or template cache."
)
# Notice the new parameter `bypass_template_cache` below
result = await self._get_prompt_impl(
prompt_name, inputs, bypass_template_cache=bypass_cache
)
self._prompt_cache.set(cache_key, result)
return result
async def get_prompt( # type: ignore
self,
name: str,
inputs: Optional[dict] = None,
prompt_override: Optional[str] = None,
) -> dict:
query = f"""
SELECT id, name, template, input_types, created_at, updated_at
FROM {self._get_table_name("prompts")}
WHERE name = $1;
"""
result = await self.connection_manager.fetchrow_query(query, [name])
if not result:
raise ValueError(f"Prompt template '{name}' not found")
input_types = result["input_types"]
if isinstance(input_types, str):
input_types = json.loads(input_types)
return {
"id": result["id"],
"name": result["name"],
"template": result["template"],
"input_types": input_types,
"created_at": result["created_at"],
"updated_at": result["updated_at"],
}
def _format_prompt(
self,
template: str,
inputs: Optional[dict[str, Any]],
input_types: dict[str, str],
) -> str:
if inputs:
# optional input validation if needed
for k, _v in inputs.items():
if k not in input_types:
raise ValueError(
f"Unexpected input '{k}' for prompt with input types {input_types}"
)
return template.format(**inputs)
return template
async def update_prompt(
self,
name: str,
template: Optional[str] = None,
input_types: Optional[dict[str, str]] = None,
) -> None:
"""Public method to update a prompt with proper cache invalidation."""
# First invalidate all caches for this prompt
self._template_cache.invalidate(name)
cache_keys_to_invalidate = [
key
for key in self._prompt_cache._cache.keys()
if key.startswith(f"{name}:") or key == name
]
for key in cache_keys_to_invalidate:
self._prompt_cache.invalidate(key)
# Perform the update
await self._update_prompt_impl(name, template, input_types)
# Force refresh template cache
template_info = await self._get_template_info(name)
if template_info:
self._template_cache.set(name, template_info)
@abstractmethod
async def _update_prompt_impl(
self,
name: str,
template: Optional[str] = None,
input_types: Optional[dict[str, str]] = None,
) -> None:
"""Implementation of prompt update logic."""
pass
@abstractmethod
async def _get_template_info(self, prompt_name: str) -> Optional[dict]:
"""Get template info with caching."""
pass
@abstractmethod
async def _get_prompt_impl(
self,
prompt_name: str,
inputs: Optional[dict[str, Any]] = None,
bypass_template_cache: bool = False,
) -> str:
"""Implementation of prompt retrieval logic."""
pass
class PostgresPromptsHandler(CacheablePromptHandler):
"""PostgreSQL implementation of the CacheablePromptHandler."""
def __init__(
self,
project_name: str,
connection_manager: PostgresConnectionManager,
prompt_directory: Optional[Path] = None,
**cache_options,
):
super().__init__(**cache_options)
self.prompt_directory = (
prompt_directory or Path(os.path.dirname(__file__)) / "prompts"
)
self.connection_manager = connection_manager
self.project_name = project_name
self.prompts: dict[str, dict[str, str | dict[str, str]]] = {}
async def _load_prompts(self) -> None:
"""Load prompts from both database and YAML files."""
# First load from database
await self._load_prompts_from_database()
# Then load from YAML files, potentially overriding unmodified database entries
await self._load_prompts_from_yaml_directory()
async def _load_prompts_from_database(self) -> None:
"""Load prompts from the database."""
query = f"""
SELECT id, name, template, input_types, created_at, updated_at
FROM {self._get_table_name("prompts")};
"""
try:
results = await self.connection_manager.fetch_query(query)
for row in results:
logger.info(f"Loading saved prompt: {row['name']}")
# Ensure input_types is a dictionary
input_types = row["input_types"]
if isinstance(input_types, str):
input_types = json.loads(input_types)
self.prompts[row["name"]] = {
"id": row["id"],
"template": row["template"],
"input_types": input_types,
"created_at": row["created_at"],
"updated_at": row["updated_at"],
}
# Pre-populate the template cache
self._template_cache.set(
row["name"],
{
"id": row["id"],
"template": row["template"],
"input_types": input_types,
},
)
logger.debug(f"Loaded {len(results)} prompts from database")
except Exception as e:
logger.error(f"Failed to load prompts from database: {e}")
raise
async def _load_prompts_from_yaml_directory(
self, default_overwrite_on_diff: bool = False
) -> None:
"""Load prompts from YAML files in the specified directory.
:param default_overwrite_on_diff: If a YAML prompt does not specify
'overwrite_on_diff', we use this default.
"""
if not self.prompt_directory.is_dir():
logger.warning(
f"Prompt directory not found: {self.prompt_directory}"
)
return
logger.info(f"Loading prompts from {self.prompt_directory}")
for yaml_file in self.prompt_directory.glob("*.yaml"):
logger.debug(f"Processing {yaml_file}")
try:
with open(yaml_file, "r", encoding="utf-8") as file:
data = yaml.safe_load(file)
if not isinstance(data, dict):
raise ValueError(
f"Invalid format in YAML file {yaml_file}"
)
for name, prompt_data in data.items():
# Attempt to parse the relevant prompt fields
template = prompt_data.get("template")
input_types = prompt_data.get("input_types", {})
# Decide on per-prompt overwrite behavior (or fallback)
overwrite_on_diff = prompt_data.get(
"overwrite_on_diff", default_overwrite_on_diff
)
# Some logic to determine if we *should* modify
# For instance, preserve only if it has never been updated
# (i.e., created_at == updated_at).
should_modify = True
if name in self.prompts:
existing = self.prompts[name]
should_modify = (
existing["created_at"]
== existing["updated_at"]
)
# If should_modify is True, the default logic is
# preserve_existing = False,
# so we can pass that in. Otherwise, preserve_existing=True
# effectively means we skip the update.
logger.info(
f"Loading default prompt: {name} from {yaml_file}."
)
await self.add_prompt(
name=name,
template=template,
input_types=input_types,
preserve_existing=False,
overwrite_on_diff=overwrite_on_diff,
)
except Exception as e:
logger.error(f"Error loading {yaml_file}: {e}")
continue
def _get_table_name(self, base_name: str) -> str:
"""Get the fully qualified table name."""
return f"{self.project_name}.{base_name}"
# Implementation of abstract methods from CacheablePromptHandler
async def _get_prompt_impl(
self,
prompt_name: str,
inputs: Optional[dict[str, Any]] = None,
bypass_template_cache: bool = False,
) -> str:
"""Implementation of database prompt retrieval."""
# If we're bypassing the template cache, skip the cache lookup
if not bypass_template_cache:
template_info = self._template_cache.get(prompt_name)
if template_info is not None:
logger.debug(f"Template cache hit: {prompt_name}")
# use that
return self._format_prompt(
template_info["template"],
inputs,
template_info["input_types"],
)
# If we get here, either no cache was found or bypass_cache is True
query = f"""
SELECT template, input_types
FROM {self._get_table_name("prompts")}
WHERE name = $1;
"""
result = await self.connection_manager.fetchrow_query(
query, [prompt_name]
)
if not result:
raise ValueError(f"Prompt template '{prompt_name}' not found")
template = result["template"]
input_types = result["input_types"]
if isinstance(input_types, str):
input_types = json.loads(input_types)
# Update template cache if not bypassing it
if not bypass_template_cache:
self._template_cache.set(
prompt_name, {"template": template, "input_types": input_types}
)
return self._format_prompt(template, inputs, input_types)
async def _get_template_info(self, prompt_name: str) -> Optional[dict]: # type: ignore
"""Get template info with caching."""
cached = self._template_cache.get(prompt_name)
if cached is not None:
return cached
query = f"""
SELECT template, input_types
FROM {self._get_table_name("prompts")}
WHERE name = $1;
"""
result = await self.connection_manager.fetchrow_query(
query, [prompt_name]
)
if result:
# Ensure input_types is a dictionary
input_types = result["input_types"]
if isinstance(input_types, str):
input_types = json.loads(input_types)
template_info = {
"template": result["template"],
"input_types": input_types,
}
self._template_cache.set(prompt_name, template_info)
return template_info
return None
async def _update_prompt_impl(
self,
name: str,
template: Optional[str] = None,
input_types: Optional[dict[str, str]] = None,
) -> None:
"""Implementation of database prompt update with proper connection
handling."""
if not template and not input_types:
return
# Clear caches first
self._template_cache.invalidate(name)
for key in list(self._prompt_cache._cache.keys()):
if key.startswith(f"{name}:"):
self._prompt_cache.invalidate(key)
# Build update query
set_clauses = []
params = [name] # First parameter is always the name
param_index = 2 # Start from 2 since $1 is name
if template:
set_clauses.append(f"template = ${param_index}")
params.append(template)
param_index += 1
if input_types:
set_clauses.append(f"input_types = ${param_index}")
params.append(json.dumps(input_types))
param_index += 1
set_clauses.append("updated_at = CURRENT_TIMESTAMP")
query = f"""
UPDATE {self._get_table_name("prompts")}
SET {", ".join(set_clauses)}
WHERE name = $1
RETURNING id, template, input_types;
"""
try:
# Execute update and get returned values
result = await self.connection_manager.fetchrow_query(
query, params
)
if not result:
raise ValueError(f"Prompt template '{name}' not found")
# Update in-memory state
if name in self.prompts:
if template:
self.prompts[name]["template"] = template
if input_types:
self.prompts[name]["input_types"] = input_types
self.prompts[name]["updated_at"] = datetime.now().isoformat()
except Exception as e:
logger.error(f"Failed to update prompt {name}: {str(e)}")
raise
async def create_tables(self):
"""Create the necessary tables for storing prompts."""
query = f"""
CREATE TABLE IF NOT EXISTS {self._get_table_name("prompts")} (
id UUID PRIMARY KEY,
name VARCHAR(255) NOT NULL UNIQUE,
template TEXT NOT NULL,
input_types JSONB NOT NULL,
created_at TIMESTAMP WITH TIME ZONE DEFAULT CURRENT_TIMESTAMP,
updated_at TIMESTAMP WITH TIME ZONE DEFAULT CURRENT_TIMESTAMP
);
CREATE OR REPLACE FUNCTION {self.project_name}.update_updated_at_column()
RETURNS TRIGGER AS $$
BEGIN
NEW.updated_at = CURRENT_TIMESTAMP;
RETURN NEW;
END;
$$ language 'plpgsql';
DROP TRIGGER IF EXISTS update_prompts_updated_at
ON {self._get_table_name("prompts")};
CREATE TRIGGER update_prompts_updated_at
BEFORE UPDATE ON {self._get_table_name("prompts")}
FOR EACH ROW
EXECUTE FUNCTION {self.project_name}.update_updated_at_column();
"""
await self.connection_manager.execute_query(query)
await self._load_prompts()
async def add_prompt(
self,
name: str,
template: str,
input_types: dict[str, str],
preserve_existing: bool = False,
overwrite_on_diff: bool = False, # <-- new param
) -> None:
"""Add or update a prompt.
If `preserve_existing` is True and prompt already exists, we skip updating.
If `overwrite_on_diff` is True and an existing prompt differs from what is provided,
we overwrite and log a warning. Otherwise, we skip if the prompt differs.
"""
# Check if prompt is in-memory
existing_prompt = self.prompts.get(name)
# If preserving existing and it already exists, skip entirely
if preserve_existing and existing_prompt:
logger.debug(
f"Preserving existing prompt: {name}, skipping update."
)
return
# If an existing prompt is found, check for diffs
if existing_prompt:
existing_template = existing_prompt["template"]
existing_input_types = existing_prompt["input_types"]
# If there's a difference in template or input_types, decide to overwrite or skip
if (
existing_template != template
or existing_input_types != input_types
):
if overwrite_on_diff:
logger.warning(
f"Overwriting existing prompt '{name}' due to detected diff."
)
else:
logger.info(
f"Prompt '{name}' differs from existing but overwrite_on_diff=False. Skipping update."
)
return
prompt_id = generate_default_prompt_id(name)
# Ensure input_types is properly serialized
input_types_json = (
json.dumps(input_types)
if isinstance(input_types, dict)
else input_types
)
# Upsert logic
query = f"""
INSERT INTO {self._get_table_name("prompts")} (id, name, template, input_types)
VALUES ($1, $2, $3, $4)
ON CONFLICT (name) DO UPDATE
SET template = EXCLUDED.template,
input_types = EXCLUDED.input_types,
updated_at = CURRENT_TIMESTAMP
RETURNING id, created_at, updated_at;
"""
result = await self.connection_manager.fetchrow_query(
query, [prompt_id, name, template, input_types_json]
)
self.prompts[name] = {
"id": result["id"],
"template": template,
"input_types": input_types,
"created_at": result["created_at"],
"updated_at": result["updated_at"],
}
# Update template cache
self._template_cache.set(
name,
{
"id": prompt_id,
"template": template,
"input_types": input_types,
},
)
# Invalidate any cached formatted prompts
for key in list(self._prompt_cache._cache.keys()):
if key.startswith(f"{name}:"):
self._prompt_cache.invalidate(key)
async def get_all_prompts(self) -> dict[str, Any]:
"""Retrieve all stored prompts."""
query = f"""
SELECT id, name, template, input_types, created_at, updated_at, COUNT(*) OVER() AS total_entries
FROM {self._get_table_name("prompts")};
"""
results = await self.connection_manager.fetch_query(query)
if not results:
return {"results": [], "total_entries": 0}
total_entries = results[0]["total_entries"] if results else 0
prompts = [
{
"name": row["name"],
"id": row["id"],
"template": row["template"],
"input_types": (
json.loads(row["input_types"])
if isinstance(row["input_types"], str)
else row["input_types"]
),
"created_at": row["created_at"],
"updated_at": row["updated_at"],
}
for row in results
]
return {"results": prompts, "total_entries": total_entries}
async def delete_prompt(self, name: str) -> None:
"""Delete a prompt template."""
query = f"""
DELETE FROM {self._get_table_name("prompts")}
WHERE name = $1;
"""
result = await self.connection_manager.execute_query(query, [name])
if result == "DELETE 0":
raise ValueError(f"Prompt template '{name}' not found")
# Invalidate caches
self._template_cache.invalidate(name)
for key in list(self._prompt_cache._cache.keys()):
if key.startswith(f"{name}:"):
self._prompt_cache.invalidate(key)
async def get_message_payload(
self,
system_prompt_name: Optional[str] = None,
system_role: str = "system",
system_inputs: dict | None = None,
system_prompt_override: Optional[str] = None,
task_prompt_name: Optional[str] = None,
task_role: str = "user",
task_inputs: Optional[dict] = None,
task_prompt: Optional[str] = None,
) -> list[dict]:
"""Create a message payload from system and task prompts."""
if system_inputs is None:
system_inputs = {}
if task_inputs is None:
task_inputs = {}
if system_prompt_override:
system_prompt = system_prompt_override
else:
system_prompt = await self.get_cached_prompt(
system_prompt_name or "system",
system_inputs,
prompt_override=system_prompt_override,
)
task_prompt = await self.get_cached_prompt(
task_prompt_name or "rag",
task_inputs,
prompt_override=task_prompt,
)
return [
{
"role": system_role,
"content": system_prompt,
},
{
"role": task_role,
"content": task_prompt,
},
]
