dimgigov
1e38e29f25
CI / test (push) Has been cancelled
CI / verify (push) Has been cancelled
Clients CI / build-server (push) Has been cancelled
Clients CI / test-python (push) Has been cancelled
Clients CI / test-javascript (push) Has been cancelled
Clients CI / test-nim (push) Has been cancelled
Clients CI / test-rust (push) Has been cancelled
- Fix SQL injection vulnerabilities in chat_history.py, langchain_store.py, rag_pipeline.py, baradb_langchain.js by switching to parameterized queries - Replace dangerous bare except: with except CatchableError:/ValueError: in llm.nim, embed.nim, cypher.nim, mcp/server.nim, executor.nim - Fix session variable leak in MCP handleVectorSearch/handleSchemaInspect - Build: 0 errors, all tests pass
236 lines
7.7 KiB
Python
236 lines
7.7 KiB
Python
"""
|
|
BaraDB Chat Message History — Conversation Buffer with RLS
|
|
|
|
Implements LangChain's BaseChatMessageHistory interface backed by BaraDB.
|
|
Supports multi-tenant isolation via tenant_id and user_id.
|
|
|
|
Usage:
|
|
from baradb import Client, WireValue
|
|
from baradb.chat_history import BaraDBChatHistory
|
|
|
|
client = Client("localhost", 9472)
|
|
await client.connect()
|
|
|
|
history = BaraDBChatHistory(
|
|
client=client,
|
|
session_id="session-123",
|
|
tenant_id="company-a",
|
|
user_id="user-42",
|
|
)
|
|
|
|
# Add messages
|
|
history.add_user_message("Hello, AI!")
|
|
history.add_ai_message("Hello, how can I help?")
|
|
|
|
# Retrieve conversation
|
|
messages = history.messages
|
|
"""
|
|
|
|
import json
|
|
from datetime import datetime
|
|
from typing import Any, Dict, List, Optional
|
|
|
|
|
|
class BaraDBChatHistory:
|
|
"""
|
|
Chat message history backed by BaraDB with multi-tenant RLS support.
|
|
|
|
Stores conversations in a `chat_history` table with columns:
|
|
id, session_id, role, content, metadata, tenant_id, user_id, created_at
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
client: Any,
|
|
session_id: str,
|
|
table: str = "chat_history",
|
|
tenant_id: Optional[str] = None,
|
|
user_id: Optional[str] = None,
|
|
max_messages: int = 1000,
|
|
):
|
|
self.client = client
|
|
self.session_id = session_id
|
|
self.table = table
|
|
self.tenant_id = tenant_id
|
|
self.user_id = user_id
|
|
self.max_messages = max_messages
|
|
self._initialized = False
|
|
|
|
async def _ensure_table(self):
|
|
if self._initialized:
|
|
return
|
|
await self.client.query(
|
|
f"""
|
|
CREATE TABLE IF NOT EXISTS {self.table} (
|
|
id TEXT PRIMARY KEY,
|
|
session_id TEXT,
|
|
role TEXT,
|
|
content TEXT,
|
|
metadata TEXT,
|
|
tenant_id TEXT,
|
|
user_id TEXT,
|
|
created_at TEXT
|
|
)
|
|
"""
|
|
)
|
|
await self.client.query(
|
|
f"CREATE INDEX IF NOT EXISTS idx_{self.table}_session "
|
|
f"ON {self.table}(session_id) USING btree"
|
|
)
|
|
self._initialized = True
|
|
|
|
def _build_session(self) -> Dict[str, str]:
|
|
s = {"app.bara_chat_session": self.session_id}
|
|
if self.tenant_id:
|
|
s["app.tenant_id"] = self.tenant_id
|
|
if self.user_id:
|
|
s["app.user_id"] = self.user_id
|
|
return s
|
|
|
|
async def add_message(self, message: Any) -> None:
|
|
await self._ensure_table()
|
|
role = getattr(message, "type", "human")
|
|
if role == "human":
|
|
role = "user"
|
|
content = getattr(message, "content", str(message))
|
|
msg_id = f"{self.session_id}:{datetime.utcnow().timestamp()}"
|
|
metadata = json.dumps(getattr(message, "additional_kwargs", {}) or {})
|
|
created_at = datetime.utcnow().isoformat()
|
|
|
|
for key, val in self._build_session().items():
|
|
await self.client.query_params(
|
|
f"SET {key} = $1", [self._wire_string(val)]
|
|
)
|
|
|
|
await self.client.query_params(
|
|
f"INSERT INTO {self.table} (id, session_id, role, content, metadata, tenant_id, user_id, created_at) "
|
|
f"VALUES ($1, $2, $3, $4, $5, $6, $7, $8)",
|
|
[
|
|
self._wire_string(msg_id),
|
|
self._wire_string(self.session_id),
|
|
self._wire_string(role),
|
|
self._wire_string(content),
|
|
self._wire_string(metadata),
|
|
self._wire_string(self.tenant_id or ""),
|
|
self._wire_string(self.user_id or ""),
|
|
self._wire_string(created_at),
|
|
],
|
|
)
|
|
|
|
def add_user_message(self, message: Any) -> None:
|
|
import asyncio
|
|
loop = asyncio.get_event_loop()
|
|
if hasattr(message, "content"):
|
|
content = message.content
|
|
else:
|
|
content = str(message)
|
|
loop.run_until_complete(self._add_message_internal(content, "user"))
|
|
|
|
def add_ai_message(self, message: Any) -> None:
|
|
import asyncio
|
|
loop = asyncio.get_event_loop()
|
|
if hasattr(message, "content"):
|
|
content = message.content
|
|
else:
|
|
content = str(message)
|
|
loop.run_until_complete(self._add_message_internal(content, "ai"))
|
|
|
|
async def _add_message_internal(self, content: str, role: str):
|
|
await self._ensure_table()
|
|
msg_id = f"{self.session_id}:{datetime.utcnow().timestamp()}"
|
|
created_at = datetime.utcnow().isoformat()
|
|
|
|
for key, val in self._build_session().items():
|
|
await self.client.query_params(
|
|
f"SET {key} = $1", [self._wire_string(val)]
|
|
)
|
|
|
|
await self.client.query_params(
|
|
f"INSERT INTO {self.table} (id, session_id, role, content, tenant_id, user_id, created_at) "
|
|
f"VALUES ($1, $2, $3, $4, $5, $6, $7)",
|
|
[
|
|
self._wire_string(msg_id),
|
|
self._wire_string(self.session_id),
|
|
self._wire_string(role),
|
|
self._wire_string(content),
|
|
self._wire_string(self.tenant_id or ""),
|
|
self._wire_string(self.user_id or ""),
|
|
self._wire_string(created_at),
|
|
],
|
|
)
|
|
|
|
async def get_messages(self) -> List[Any]:
|
|
await self._ensure_table()
|
|
class SimpleMessage:
|
|
def __init__(self, role: str, content: str):
|
|
self.type = "human" if role == "user" else role
|
|
self.content = content
|
|
self.additional_kwargs = {}
|
|
|
|
def __repr__(self):
|
|
return f"{self.type}: {self.content}"
|
|
|
|
for key, val in self._build_session().items():
|
|
await self.client.query_params(
|
|
f"SET {key} = $1", [self._wire_string(val)]
|
|
)
|
|
|
|
result = await self.client.query_params(
|
|
f"SELECT role, content FROM {self.table} "
|
|
f"WHERE session_id = $1 "
|
|
f"ORDER BY created_at ASC "
|
|
f"LIMIT $2",
|
|
[
|
|
self._wire_string(self.session_id),
|
|
self._wire_int(self.max_messages),
|
|
],
|
|
)
|
|
messages = []
|
|
if result and hasattr(result, "rows"):
|
|
for row in result.rows:
|
|
role = row.get("role", "user")
|
|
content = row.get("content", "")
|
|
messages.append(SimpleMessage(role, content))
|
|
|
|
return messages
|
|
|
|
@property
|
|
def messages(self) -> List[Any]:
|
|
import asyncio
|
|
loop = asyncio.get_event_loop()
|
|
return loop.run_until_complete(self.get_messages())
|
|
|
|
async def clear(self) -> None:
|
|
await self._ensure_table()
|
|
for key, val in self._build_session().items():
|
|
await self.client.query_params(
|
|
f"SET {key} = $1", [self._wire_string(val)]
|
|
)
|
|
await self.client.query_params(
|
|
f"DELETE FROM {self.table} WHERE session_id = $1",
|
|
[self._wire_string(self.session_id)],
|
|
)
|
|
|
|
async def get_session_summary(self, max_tokens: int = 2000) -> str:
|
|
messages = await self.get_messages()
|
|
parts = []
|
|
total_chars = 0
|
|
for msg in reversed(messages):
|
|
text = f"{msg.type}: {getattr(msg, 'content', '')}"
|
|
if total_chars + len(text) > max_tokens * 4:
|
|
break
|
|
parts.insert(0, text)
|
|
total_chars += len(text)
|
|
return "\n".join(parts)
|
|
|
|
@staticmethod
|
|
def _wire_string(val: str) -> Any:
|
|
# Lazy import to avoid circular dependency
|
|
from baradb import WireValue
|
|
return WireValue.string(val)
|
|
|
|
@staticmethod
|
|
def _wire_int(val: int) -> Any:
|
|
from baradb import WireValue
|
|
return WireValue.int64(val)
|