JDSS/mem0
2
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix: sql injection, prompt injection (#4997)

Browse Source 
Co-authored-by: kartik-mem0 <kartik.labhshetwar@mem0.ai>
This commit is contained in:
Harsh Vardhan Gupta
2026-04-29 00:51:16 +05:30
committed by GitHub
co-authored by kartik-mem0
parent b66cf0f272
commit 1b95c99db4
8 changed files with 233 additions and 150 deletions
Download Patch File Download Diff File Expand all files Collapse all files
mem0/reranker/llm_reranker.py
+37 -21
View File
@@ -58,24 +58,35 @@ class LLMReranker(BaseReranker):
# Initialize LLM using the factory
self.llm = LlmFactory.create(llm_provider, llm_config)
# Default scoring prompt
self.scoring_prompt = getattr(self.config, 'scoring_prompt', None) or self._get_default_prompt()
def _get_default_prompt(self) -> str:
"""Get the default scoring prompt template."""
return """You are a relevance scoring assistant. Given a query and a document, you need to score how relevant the document is to the query.
# Honor custom scoring_prompt from config if provided
custom_prompt = getattr(self.config, 'scoring_prompt', None)
if custom_prompt:
import warnings
warnings.warn(
"LLMRerankerConfig.scoring_prompt is deprecated and will be removed in a future version. "
"The prompt is now used as the system message.",
DeprecationWarning,
stacklevel=2,
)
self._system_prompt = custom_prompt
else:
self._system_prompt = self._SYSTEM_PROMPT
Score the relevance on a scale from 0.0 to 1.0, where:
- 1.0 = Perfectly relevant and directly answers the query
- 0.8-0.9 = Highly relevant with good information
- 0.6-0.7 = Moderately relevant with some useful information
- 0.4-0.5 = Slightly relevant with limited useful information
- 0.0-0.3 = Not relevant or no useful information
_SYSTEM_PROMPT = (
"You are a relevance scoring assistant. "
"Given a query and a document, score how relevant the document is to the query.\n\n"
"Score the relevance on a scale from 0.0 to 1.0, where:\n"
"- 1.0 = Perfectly relevant and directly answers the query\n"
"- 0.8-0.9 = Highly relevant with good information\n"
"- 0.6-0.7 = Moderately relevant with some useful information\n"
"- 0.4-0.5 = Slightly relevant with limited useful information\n"
"- 0.0-0.3 = Not relevant or no useful information\n\n"
"Respond with only a single numerical score between 0.0 and 1.0. "
"Do not include any explanation or additional text."
)
Query: "{query}"
Document: "{document}"
Provide only a single numerical score between 0.0 and 1.0. Do not include any explanation or additional text."""
# Maximum character length for query and document inputs to prevent prompt flooding.
_MAX_INPUT_LEN = 4000
def _extract_score(self, response_text: str) -> float:
"""Extract numerical score from LLM response."""
@@ -119,12 +130,17 @@ Provide only a single numerical score between 0.0 and 1.0. Do not include any ex
doc_text = str(doc)
try:
# Generate scoring prompt
prompt = self.scoring_prompt.format(query=query, document=doc_text)
# Get LLM response
# Truncate inputs to prevent prompt flooding, then send as separate
# system/user messages so instructions cannot be overridden by user data.
safe_query = query[: self._MAX_INPUT_LEN]
safe_doc = doc_text[: self._MAX_INPUT_LEN]
user_message = f"Query: {safe_query}\n\nDocument: {safe_doc}"
response = self.llm.generate_response(
messages=[{"role": "user", "content": prompt}]
messages=[
{"role": "system", "content": self._system_prompt},
{"role": "user", "content": user_message},
]
)
# Extract score from response
mem0/vector_stores/azure_mysql.py
+14 -2
View File
@@ -1,5 +1,6 @@
import json
import logging
import re
from contextlib import contextmanager
from typing import Any, Dict, List, Optional
@@ -25,6 +26,17 @@ from mem0.vector_stores.base import VectorStoreBase
logger = logging.getLogger(__name__)
_SAFE_IDENTIFIER_RE = re.compile(r'^[a-zA-Z_][a-zA-Z0-9_]{0,127}$')
def _validate_identifier(name: str, label: str = "identifier") -> str:
if not _SAFE_IDENTIFIER_RE.match(name):
raise ValueError(
f"Invalid {label} '{name}': only letters, digits, and underscores are allowed, "
"must start with a letter or underscore, and be at most 128 characters."
)
return name
class OutputData(BaseModel):
id: Optional[str]
@@ -72,7 +84,7 @@ class AzureMySQL(VectorStoreBase):
self.user = user
self.password = password
self.database = database
self.collection_name = collection_name
self.collection_name = _validate_identifier(collection_name, "collection_name")
self.embedding_model_dims = embedding_model_dims
self.use_azure_credential = use_azure_credential
self.ssl_ca = ssl_ca
@@ -174,7 +186,7 @@ class AzureMySQL(VectorStoreBase):
vector_size (int, optional): Vector dimension (uses self.embedding_model_dims if not provided)
distance (str): Distance metric (cosine, euclidean, dot_product)
"""
table_name = name or self.collection_name
table_name = _validate_identifier(name, "table_name") if name else self.collection_name
dims = vector_size or self.embedding_model_dims
with self._get_cursor(commit=True) as cur:
mem0/vector_stores/cassandra.py
+19 -9
View File
@@ -1,5 +1,6 @@
import json
import logging
import re
import uuid
from typing import Any, Dict, List, Optional
@@ -19,6 +20,17 @@ from mem0.vector_stores.base import VectorStoreBase
logger = logging.getLogger(__name__)
_SAFE_IDENTIFIER_RE = re.compile(r'^[a-zA-Z_][a-zA-Z0-9_]{0,127}$')
def _validate_identifier(name: str, label: str = "identifier") -> str:
if not _SAFE_IDENTIFIER_RE.match(name):
raise ValueError(
f"Invalid {label} '{name}': only letters, digits, and underscores are allowed, "
"must start with a letter or underscore, and be at most 128 characters."
)
return name
class OutputData(BaseModel):
id: Optional[str]
@@ -59,8 +71,8 @@ class CassandraDB(VectorStoreBase):
self.port = port
self.username = username
self.password = password
self.keyspace = keyspace
self.collection_name = collection_name
self.keyspace = _validate_identifier(keyspace, "keyspace")
self.collection_name = _validate_identifier(collection_name, "collection_name")
self.embedding_model_dims = embedding_model_dims
self.secure_connect_bundle = secure_connect_bundle
self.protocol_version = protocol_version
@@ -156,7 +168,7 @@ class CassandraDB(VectorStoreBase):
vector_size (int, optional): Vector dimension (uses self.embedding_model_dims if not provided)
distance (str): Distance metric (cosine, euclidean, dot_product)
"""
table_name = name or self.collection_name
table_name = _validate_identifier(name, "table_name") if name else self.collection_name
dims = vector_size or self.embedding_model_dims
try:
@@ -375,12 +387,10 @@ class CassandraDB(VectorStoreBase):
List[str]: List of collection names
"""
try:
query = f"""
SELECT table_name
FROM system_schema.tables
WHERE keyspace_name = '{self.keyspace}'
"""
rows = self.session.execute(query)
prepared = self.session.prepare(
"SELECT table_name FROM system_schema.tables WHERE keyspace_name = ?"
)
rows = self.session.execute(prepared, (self.keyspace,))
return [row.table_name for row in rows]
except Exception as e:
logger.error(f"Failed to list collections: {e}")
mem0/vector_stores/pgvector.py
+60 -49
View File
@@ -7,6 +7,7 @@ from pydantic import BaseModel
# Try to import psycopg (psycopg3) first, then fall back to psycopg2
try:
from psycopg import sql
from psycopg.types.json import Json
from psycopg_pool import ConnectionPool
PSYCOPG_VERSION = 3
@@ -14,6 +15,7 @@ try:
logger.info("Using psycopg (psycopg3) with ConnectionPool for PostgreSQL connections")
except ImportError:
try:
from psycopg2 import sql
from psycopg2.extras import Json, execute_values
from psycopg2.pool import ThreadedConnectionPool as ConnectionPool
PSYCOPG_VERSION = 2
@@ -144,6 +146,10 @@ class PGVector(VectorStoreBase):
cur.close()
self.connection_pool.putconn(conn)
def _col(self) -> "sql.Identifier":
"""Return a safely-quoted SQL identifier for the collection table."""
return sql.Identifier(self.collection_name)
def create_col(self) -> None:
"""
Create a new collection (table in PostgreSQL).
@@ -152,39 +158,45 @@ class PGVector(VectorStoreBase):
with self._get_cursor(commit=True) as cur:
cur.execute("CREATE EXTENSION IF NOT EXISTS vector")
cur.execute(
f"""
CREATE TABLE IF NOT EXISTS {self.collection_name} (
sql.SQL("""
CREATE TABLE IF NOT EXISTS {} (
id UUID PRIMARY KEY,
vector vector({self.embedding_model_dims}),
vector vector({}),
payload JSONB
);
"""
""").format(self._col(), sql.Literal(self.embedding_model_dims))
)
if self.use_diskann and self.embedding_model_dims < 2000:
cur.execute("SELECT * FROM pg_extension WHERE extname = 'vectorscale'")
if cur.fetchone():
# Create DiskANN index if extension is installed for faster search
cur.execute(
f"""
CREATE INDEX IF NOT EXISTS {self.collection_name}_diskann_idx
ON {self.collection_name}
sql.SQL("""
CREATE INDEX IF NOT EXISTS {} ON {}
USING diskann (vector);
"""
""").format(
sql.Identifier(f"{self.collection_name}_diskann_idx"),
self._col(),
)
)
elif self.use_hnsw:
cur.execute(
f"""
CREATE INDEX IF NOT EXISTS {self.collection_name}_hnsw_idx
ON {self.collection_name}
sql.SQL("""
CREATE INDEX IF NOT EXISTS {} ON {}
USING hnsw (vector vector_cosine_ops)
"""
""").format(
sql.Identifier(f"{self.collection_name}_hnsw_idx"),
self._col(),
)
)
cur.execute(
f"""
CREATE INDEX IF NOT EXISTS {self.collection_name}_text_lemmatized_idx
ON {self.collection_name}
sql.SQL("""
CREATE INDEX IF NOT EXISTS {} ON {}
USING gin(to_tsvector('simple', payload->>'text_lemmatized'));
"""
""").format(
sql.Identifier(f"{self.collection_name}_text_lemmatized_idx"),
self._col(),
)
)
def insert(self, vectors: list[list[float]], payloads=None, ids=None) -> None:
@@ -195,14 +207,14 @@ class PGVector(VectorStoreBase):
if PSYCOPG_VERSION == 3:
with self._get_cursor(commit=True) as cur:
cur.executemany(
f"INSERT INTO {self.collection_name} (id, vector, payload) VALUES (%s, %s, %s)",
sql.SQL("INSERT INTO {} (id, vector, payload) VALUES (%s, %s, %s)").format(self._col()),
data,
)
else:
with self._get_cursor(commit=True) as cur:
execute_values(
cur,
f"INSERT INTO {self.collection_name} (id, vector, payload) VALUES %s",
sql.SQL("INSERT INTO {} (id, vector, payload) VALUES %s").format(self._col()),
data,
)
@@ -233,17 +245,17 @@ class PGVector(VectorStoreBase):
filter_conditions.append("payload->>%s = %s")
filter_params.extend([k, str(v)])
filter_clause = "WHERE " + " AND ".join(filter_conditions) if filter_conditions else ""
filter_clause = sql.SQL("WHERE " + " AND ".join(filter_conditions)) if filter_conditions else sql.SQL("")
with self._get_cursor() as cur:
cur.execute(
f"""
sql.SQL("""
SELECT id, vector <=> %s::vector AS distance, payload
FROM {self.collection_name}
{filter_clause}
FROM {}
{}
ORDER BY distance
LIMIT %s
""",
""").format(self._col(), filter_clause),
(vectors, *filter_params, top_k),
)
@@ -270,21 +282,19 @@ class PGVector(VectorStoreBase):
filter_conditions.append("payload->>%s = %s")
filter_params.extend([k, str(v)])
filter_clause = ""
if filter_conditions:
filter_clause = "AND " + " AND ".join(filter_conditions)
filter_clause = sql.SQL("AND " + " AND ".join(filter_conditions)) if filter_conditions else sql.SQL("")
try:
with self._get_cursor() as cur:
cur.execute(
f"""
sql.SQL("""
SELECT id, ts_rank_cd(to_tsvector('simple', payload->>'text_lemmatized'), plainto_tsquery('simple', %s)) AS score, payload
FROM {self.collection_name}
FROM {}
WHERE to_tsvector('simple', payload->>'text_lemmatized') @@ plainto_tsquery('simple', %s)
{filter_clause}
{}
ORDER BY score DESC
LIMIT %s
""",
""").format(self._col(), filter_clause),
(query, query, *filter_params, top_k),
)
@@ -302,7 +312,7 @@ class PGVector(VectorStoreBase):
vector_id (str): ID of the vector to delete.
"""
with self._get_cursor(commit=True) as cur:
cur.execute(f"DELETE FROM {self.collection_name} WHERE id = %s", (vector_id,))
cur.execute(sql.SQL("DELETE FROM {} WHERE id = %s").format(self._col()), (vector_id,))
def update(
self,
@@ -321,7 +331,7 @@ class PGVector(VectorStoreBase):
with self._get_cursor(commit=True) as cur:
if vector:
cur.execute(
f"UPDATE {self.collection_name} SET vector = %s WHERE id = %s",
sql.SQL("UPDATE {} SET vector = %s WHERE id = %s").format(self._col()),
(vector, vector_id),
)
if payload:
@@ -329,13 +339,13 @@ class PGVector(VectorStoreBase):
if PSYCOPG_VERSION == 3:
# psycopg3 uses psycopg.types.json.Json
cur.execute(
f"UPDATE {self.collection_name} SET payload = %s WHERE id = %s",
sql.SQL("UPDATE {} SET payload = %s WHERE id = %s").format(self._col()),
(Json(payload), vector_id),
)
else:
# psycopg2 uses psycopg2.extras.Json
cur.execute(
f"UPDATE {self.collection_name} SET payload = %s WHERE id = %s",
sql.SQL("UPDATE {} SET payload = %s WHERE id = %s").format(self._col()),
(Json(payload), vector_id),
)
@@ -352,7 +362,7 @@ class PGVector(VectorStoreBase):
"""
with self._get_cursor() as cur:
cur.execute(
f"SELECT id, vector, payload FROM {self.collection_name} WHERE id = %s",
sql.SQL("SELECT id, vector, payload FROM {} WHERE id = %s").format(self._col()),
(vector_id,),
)
result = cur.fetchone()
@@ -374,7 +384,7 @@ class PGVector(VectorStoreBase):
def delete_col(self) -> None:
"""Delete a collection."""
with self._get_cursor(commit=True) as cur:
cur.execute(f"DROP TABLE IF EXISTS {self.collection_name}")
cur.execute(sql.SQL("DROP TABLE IF EXISTS {}").format(self._col()))
def col_info(self) -> dict[str, Any]:
"""
@@ -385,14 +395,14 @@ class PGVector(VectorStoreBase):
"""
with self._get_cursor() as cur:
cur.execute(
f"""
sql.SQL("""
SELECT
table_name,
(SELECT COUNT(*) FROM {self.collection_name}) as row_count,
(SELECT pg_size_pretty(pg_total_relation_size('{self.collection_name}'))) as total_size
(SELECT COUNT(*) FROM {}) as row_count,
(SELECT pg_size_pretty(pg_total_relation_size({}::regclass))) as total_size
FROM information_schema.tables
WHERE table_schema = 'public' AND table_name = %s
""",
""").format(self._col(), sql.Literal(self.collection_name)),
(self.collection_name,),
)
result = cur.fetchone()
@@ -421,17 +431,18 @@ class PGVector(VectorStoreBase):
filter_conditions.append("payload->>%s = %s")
filter_params.extend([k, str(v)])
filter_clause = "WHERE " + " AND ".join(filter_conditions) if filter_conditions else ""
query = f"""
SELECT id, vector, payload
FROM {self.collection_name}
{filter_clause}
LIMIT %s
"""
filter_clause = sql.SQL("WHERE " + " AND ".join(filter_conditions)) if filter_conditions else sql.SQL("")
with self._get_cursor() as cur:
cur.execute(query, (*filter_params, top_k))
cur.execute(
sql.SQL("""
SELECT id, vector, payload
FROM {}
{}
LIMIT %s
""").format(self._col(), filter_clause),
(*filter_params, top_k),
)
results = cur.fetchall()
return [[OutputData(id=str(r[0]), score=None, payload=r[2]) for r in results]]