feat: pg migration, more tables

This commit is contained in:
Alex committed 2026-04-12 12:15:59 +01:00
1 parent cada1a44fc
commit 502819ae52
36 files changed
+1929 -87

No files matched your search

@@ -1,51 +1,80 @@
Ollama
Qdrant
Milvus
Chatwoot
Nextra
VSCode
npm
LLMs
Agentic
Anthropic's
api
APIs
Groq
SGLang
LMDeploy
OAuth
Vite
LLM
JSONPath
UIs
configs
uncomment
qdrant
vectorstore
docsgpt
llm
GPUs
kubectl
Lightsail
enqueues
chatbot
VSCode's
Shareability
feedbacks
Atlassian
automations
Premade
Signup
Repo
repo
env
URl
agentic
llama_cpp
parsable
SDKs
boolean
bool
hardcode
EOL
Postgres
Supabase
config
autoescaping
Autoescaping
backfill
backfills
bool
boolean
brave_web_search
chatbot
Chatwoot
config
configs
CSVs
dev
diarization
Docling
docsgpt
docstrings
Entra
env
enqueues
EOL
ESLint
feedbacks
Figma
GPUs
Groq
hardcode
hardcoding
Idempotency
JSONPath
kubectl
Lightsail
llama_cpp
llm
LLM
LLMs
LMDeploy
Milvus
Mixtral
namespace
namespaces
needs_auth
Nextra
Novita
npm
OAuth
Ollama
opencode
parsable
passthrough
PDFs
pgvector
Postgres
Premade
Pydantic
pytest
Qdrant
qdrant
Repo
repo
Sanitization
SDKs
SGLang
Shareability
Signup
Supabase
UIs
uncomment
URl
vectorstore
Vite
VSCode
VSCode's
widget's
+2 -2
View File
@@ -73,7 +73,7 @@ class BraveSearchTool(Tool):
"X-Subscription-Token": self.token,
}
response = requests.get(url, params=params, headers=headers)
response = requests.get(url, params=params, headers=headers, timeout=100)
if response.status_code == 200:
return {
@@ -118,7 +118,7 @@ class BraveSearchTool(Tool):
"X-Subscription-Token": self.token,
}
response = requests.get(url, params=params, headers=headers)
response = requests.get(url, params=params, headers=headers, timeout=100)
if response.status_code == 200:
return {
+1 -1
View File
@@ -28,7 +28,7 @@ class CryptoPriceTool(Tool):
returns price in USD.
"""
url = f"https://min-api.cryptocompare.com/data/price?fsym={symbol.upper()}&tsyms={currency.upper()}"
response = requests.get(url)
response = requests.get(url, timeout=100)
if response.status_code == 200:
data = response.json()
if currency.upper() in data:
+1 -1
View File
@@ -71,7 +71,7 @@ class NtfyTool(Tool):
if self.token:
headers["Authorization"] = f"Basic {self.token}"
data = message.encode("utf-8")
response = requests.post(url, headers=headers, data=data)
response = requests.post(url, headers=headers, data=data, timeout=100)
return {"status_code": response.status_code, "message": "Message sent"}
def get_actions_metadata(self):
+2 -2
View File
@@ -31,14 +31,14 @@ class TelegramTool(Tool):
logger.debug("Sending Telegram message to chat_id=%s", chat_id)
url = f"https://api.telegram.org/bot{self.token}/sendMessage"
payload = {"chat_id": chat_id, "text": text}
response = requests.post(url, data=payload)
response = requests.post(url, data=payload, timeout=100)
return {"status_code": response.status_code, "message": "Message sent"}
def _send_image(self, image_url, chat_id):
logger.debug("Sending Telegram image to chat_id=%s", chat_id)
url = f"https://api.telegram.org/bot{self.token}/sendPhoto"
payload = {"chat_id": chat_id, "photo": image_url}
response = requests.post(url, data=payload)
response = requests.post(url, data=payload, timeout=100)
return {"status_code": response.status_code, "message": "Image sent"}
def get_actions_metadata(self):
+12
View File
@@ -469,6 +469,18 @@ class BaseAnswerResource:
log_data[key] = value[:10000]
self.user_logs_collection.insert_one(log_data)
from application.storage.db.dual_write import dual_write
from application.storage.db.repositories.user_logs import UserLogsRepository
dual_write(
UserLogsRepository,
lambda repo, d=log_data: repo.insert(
user_id=d.get("user"),
endpoint="stream_answer",
data=d,
),
)
data = json.dumps({"type": "end"})
yield f"data: {data}\n\n"
except GeneratorExit:
+14
View File
@@ -8,6 +8,8 @@ from flask_restx import fields, Namespace, Resource
from application.api import api
from application.api.user.base import current_dir, prompts_collection
from application.storage.db.dual_write import dual_write
from application.storage.db.repositories.prompts import PromptsRepository
from application.utils import check_required_fields
prompts_ns = Namespace(
@@ -49,6 +51,10 @@ class CreatePrompt(Resource):
}
)
new_id = str(resp.inserted_id)
dual_write(
PromptsRepository,
lambda repo, u=user, n=data["name"], c=data["content"]: repo.create(u, n, c),
)
except Exception as err:
current_app.logger.error(f"Error creating prompt: {err}", exc_info=True)
return make_response(jsonify({"success": False}), 400)
@@ -149,6 +155,10 @@ class DeletePrompt(Resource):
return missing_fields
try:
prompts_collection.delete_one({"_id": ObjectId(data["id"]), "user": user})
dual_write(
PromptsRepository,
lambda repo, pid=data["id"], u=user: repo.delete(pid, u),
)
except Exception as err:
current_app.logger.error(f"Error deleting prompt: {err}", exc_info=True)
return make_response(jsonify({"success": False}), 400)
@@ -185,6 +195,10 @@ class UpdatePrompt(Resource):
{"_id": ObjectId(data["id"]), "user": user},
{"$set": {"name": data["name"], "content": data["content"]}},
)
dual_write(
PromptsRepository,
lambda repo, pid=data["id"], u=user, n=data["name"], c=data["content"]: repo.update(pid, u, n, c),
)
except Exception as err:
current_app.logger.error(f"Error updating prompt: {err}", exc_info=True)
return make_response(jsonify({"success": False}), 400)
+13
View File
@@ -9,6 +9,8 @@ from application.agents.tools.tool_manager import ToolManager
from application.api import api
from application.api.user.base import user_tools_collection
from application.core.url_validation import SSRFError, validate_url
from application.storage.db.dual_write import dual_write
from application.storage.db.repositories.user_tools import UserToolsRepository
from application.security.encryption import decrypt_credentials, encrypt_credentials
from application.utils import check_required_fields, validate_function_name
@@ -294,6 +296,13 @@ class CreateTool(Resource):
}
resp = user_tools_collection.insert_one(new_tool)
new_id = str(resp.inserted_id)
dual_write(
UserToolsRepository,
lambda repo, u=user, t=new_tool: repo.create(
u, t["name"], config=t.get("config"),
custom_name=t.get("customName"), display_name=t.get("displayName"),
),
)
except Exception as err:
current_app.logger.error(f"Error creating tool: {err}", exc_info=True)
return make_response(jsonify({"success": False}), 400)
@@ -581,6 +590,10 @@ class DeleteTool(Resource):
result = user_tools_collection.delete_one(
{"_id": ObjectId(data["id"]), "user": user}
)
dual_write(
UserToolsRepository,
lambda repo, tid=data["id"], u=user: repo.delete(tid, u),
)
if result.deleted_count == 0:
return make_response(
jsonify({"success": False, "message": "Tool not found"}), 404
+4 -4
View File
@@ -28,13 +28,13 @@ class Settings(BaseSettings):
CELERY_RESULT_BACKEND: str = "redis://localhost:6379/1"
MONGO_URI: str = "mongodb://localhost:27017/docsgpt"
MONGO_DB_NAME: str = "docsgpt"
# User-data Postgres DB Optional during the MongoDB→Postgres migration; becomes required once the migration is
# complete.
# User-data Postgres DB. Falls back to PGVECTOR_CONNECTION_STRING when
# unset, so operators sharing a single Postgres cluster for both vectors
# and app data don't need a second env var.
POSTGRES_URI: Optional[str] = None
# MongoDB→Postgres migration switches
# MongoDB→Postgres migration: dual-write to Postgres (Mongo stays source of truth)
USE_POSTGRES: bool = False
READ_POSTGRES: bool = False
LLM_PATH: str = os.path.join(current_dir, "models/docsgpt-7b-f16.gguf")
DEFAULT_MAX_HISTORY: int = 150
DEFAULT_LLM_TOKEN_LIMIT: int = 128000 # Fallback when model not found in registry
+16
View File
@@ -157,5 +157,21 @@ def _log_to_mongodb(
user_logs_collection.insert_one(log_entry)
logging.debug(f"Logged activity to MongoDB: {activity_id}")
from application.storage.db.dual_write import dual_write
from application.storage.db.repositories.stack_logs import StackLogsRepository
dual_write(
StackLogsRepository,
lambda repo, e=log_entry: repo.insert(
activity_id=e["id"],
endpoint=e.get("endpoint"),
level=e.get("level"),
user_id=e.get("user"),
api_key=e.get("api_key"),
query=e.get("query"),
stacks=e.get("stacks"),
),
)
except Exception as e:
logging.error(f"Failed to log to MongoDB: {e}", exc_info=True)
@@ -205,7 +205,7 @@ class SharePointLoader(BaseConnectorLoader):
try:
url = self._get_item_url(file_id)
params = {'$select': 'id,name,file,createdDateTime,lastModifiedDateTime,size'}
response = requests.get(url, headers=self._get_headers(), params=params)
response = requests.get(url, headers=self._get_headers(), params=params, timeout=100)
response.raise_for_status()
file_metadata = response.json()
@@ -236,9 +236,9 @@ class SharePointLoader(BaseConnectorLoader):
search_url = f"{self.GRAPH_API_BASE}/drives/{drive_id}/root/search(q='{encoded_query}')"
else:
search_url = f"{self.GRAPH_API_BASE}/me/drive/search(q='{encoded_query}')"
response = requests.get(search_url, headers=self._get_headers(), params=params)
response = requests.get(search_url, headers=self._get_headers(), params=params, timeout=100)
else:
response = requests.get(url, headers=self._get_headers(), params=params)
response = requests.get(url, headers=self._get_headers(), params=params, timeout=100)
response.raise_for_status()
@@ -307,7 +307,8 @@ class SharePointLoader(BaseConnectorLoader):
response = requests.get(
f"{self.GRAPH_API_BASE}/me/drive",
headers=self._get_headers(),
params={'$select': 'webUrl'}
params={'$select': 'webUrl'},
timeout=100,
)
response.raise_for_status()
return response.json().get('webUrl')
@@ -352,7 +353,7 @@ class SharePointLoader(BaseConnectorLoader):
headers = self._get_headers()
headers["Content-Type"] = "application/json"
response = requests.post(url, headers=headers, json=body)
response = requests.post(url, headers=headers, json=body, timeout=100)
response.raise_for_status()
results = response.json()
@@ -472,7 +473,7 @@ class SharePointLoader(BaseConnectorLoader):
try:
url = f"{self._get_item_url(file_id)}/content"
response = requests.get(url, headers=self._get_headers())
response = requests.get(url, headers=self._get_headers(), timeout=100)
response.raise_for_status()
try:
@@ -491,7 +492,7 @@ class SharePointLoader(BaseConnectorLoader):
try:
url = self._get_item_url(file_id)
params = {'$select': 'id,name,file'}
response = requests.get(url, headers=self._get_headers(), params=params)
response = requests.get(url, headers=self._get_headers(), params=params, timeout=100)
response.raise_for_status()
metadata = response.json()
@@ -507,7 +508,7 @@ class SharePointLoader(BaseConnectorLoader):
full_path = os.path.join(local_dir, file_name)
download_url = f"{self._get_item_url(file_id)}/content"
download_response = requests.get(download_url, headers=self._get_headers())
download_response = requests.get(download_url, headers=self._get_headers(), timeout=100)
download_response.raise_for_status()
with open(full_path, 'wb') as f:
@@ -527,7 +528,7 @@ class SharePointLoader(BaseConnectorLoader):
params = {'$top': 1000}
while url:
response = requests.get(url, headers=self._get_headers(), params=params)
response = requests.get(url, headers=self._get_headers(), params=params, timeout=100)
response.raise_for_status()
results = response.json()
@@ -609,7 +610,7 @@ class SharePointLoader(BaseConnectorLoader):
try:
url = self._get_item_url(folder_id)
params = {'$select': 'id,name'}
response = requests.get(url, headers=self._get_headers(), params=params)
response = requests.get(url, headers=self._get_headers(), params=params, timeout=100)
response.raise_for_status()
folder_metadata = response.json()
+1 -1
View File
@@ -24,7 +24,7 @@ class PDFParser(BaseParser):
# alternatively you can use local vision capable LLM
with open(file, "rb") as file_loaded:
files = {'file': file_loaded}
response = requests.post(doc2md_service, files=files)
response = requests.post(doc2md_service, files=files, timeout=100)
data = response.json()["markdown"]
return data
+1 -1
View File
@@ -24,7 +24,7 @@ class ImageParser(BaseParser):
# alternatively you can use local vision capable LLM
with open(file, "rb") as file_loaded:
files = {'file': file_loaded}
response = requests.post(doc2md_service, files=files)
response = requests.post(doc2md_service, files=files, timeout=100)
data = response.json()["markdown"]
else:
data = ""
+1 -1
View File
@@ -77,7 +77,7 @@ class GitHubLoader(BaseRemote):
def _make_request(self, url: str, max_retries: int = 3) -> requests.Response:
"""Make a request with retry logic for rate limiting"""
for attempt in range(max_retries):
response = requests.get(url, headers=self.headers)
response = requests.get(url, headers=self.headers, timeout=100)
if response.status_code == 200:
return response
+30 -11
View File
@@ -23,13 +23,38 @@ from application.core.settings import settings
_engine: Optional[Engine] = None
def _resolve_uri() -> str:
"""Pick the Postgres URI, falling back to the pgvector connection string.
If ``POSTGRES_URI`` is set explicitly, use it. Otherwise, if the app
is already using pgvector with a Postgres-backed vector store
(``PGVECTOR_CONNECTION_STRING``), reuse that same cluster for
user-data tables — they can share a database. The pgvector URI is
normalized from its libpq form to the SQLAlchemy psycopg3 dialect.
"""
if settings.POSTGRES_URI:
return settings.POSTGRES_URI
if settings.PGVECTOR_CONNECTION_STRING:
from application.core.db_uri import normalize_postgres_uri
uri = normalize_postgres_uri(settings.PGVECTOR_CONNECTION_STRING)
if uri:
return uri
raise RuntimeError(
"POSTGRES_URI is not configured and no PGVECTOR_CONNECTION_STRING "
"to fall back to. Set POSTGRES_URI in your .env to a URI such as "
"'postgresql+psycopg://user:pass@host:5432/docsgpt'."
)
def get_engine() -> Engine:
"""Return the process-wide SQLAlchemy Engine, creating it if needed.
Raises:
RuntimeError: If ``settings.POSTGRES_URI`` is unset. Callers that
reach this path without a configured URI have a setup bug — the
error message points them at the right setting.
Falls back to ``PGVECTOR_CONNECTION_STRING`` when ``POSTGRES_URI``
is not set, so operators using pgvector on the same cluster don't
need a second env var.
Returns:
A SQLAlchemy ``Engine`` configured with a pooled connection to
@@ -37,14 +62,8 @@ def get_engine() -> Engine:
"""
global _engine
if _engine is None:
if not settings.POSTGRES_URI:
raise RuntimeError(
"POSTGRES_URI is not configured. Set it in your .env to a "
"psycopg3 URI such as "
"'postgresql+psycopg://user:pass@host:5432/docsgpt'."
)
_engine = create_engine(
settings.POSTGRES_URI,
_resolve_uri(),
pool_size=10,
max_overflow=20,
pool_pre_ping=True, # survive PgBouncer / idle-disconnect recycles
+73
View File
@@ -8,8 +8,10 @@ If the two drift, migrations win — update this file to match.
"""
from sqlalchemy import (
BigInteger,
Column,
DateTime,
Integer,
MetaData,
Table,
Text,
@@ -36,3 +38,74 @@ users_table = Table(
Column("created_at", DateTime(timezone=True), nullable=False, server_default=func.now()),
Column("updated_at", DateTime(timezone=True), nullable=False, server_default=func.now()),
)
prompts_table = Table(
"prompts",
metadata,
Column("id", UUID(as_uuid=True), primary_key=True, server_default=func.gen_random_uuid()),
Column("user_id", Text, nullable=False),
Column("name", Text, nullable=False),
Column("content", Text, nullable=False),
Column("created_at", DateTime(timezone=True), nullable=False, server_default=func.now()),
Column("updated_at", DateTime(timezone=True), nullable=False, server_default=func.now()),
)
user_tools_table = Table(
"user_tools",
metadata,
Column("id", UUID(as_uuid=True), primary_key=True, server_default=func.gen_random_uuid()),
Column("user_id", Text, nullable=False),
Column("name", Text, nullable=False),
Column("custom_name", Text),
Column("display_name", Text),
Column("config", JSONB, nullable=False, server_default="{}"),
Column("created_at", DateTime(timezone=True), nullable=False, server_default=func.now()),
Column("updated_at", DateTime(timezone=True), nullable=False, server_default=func.now()),
)
token_usage_table = Table(
"token_usage",
metadata,
Column("id", BigInteger, primary_key=True, autoincrement=True),
Column("user_id", Text),
Column("api_key", Text),
Column("agent_id", UUID(as_uuid=True)),
Column("prompt_tokens", Integer, nullable=False, server_default="0"),
Column("generated_tokens", Integer, nullable=False, server_default="0"),
Column("timestamp", DateTime(timezone=True), nullable=False, server_default=func.now()),
)
user_logs_table = Table(
"user_logs",
metadata,
Column("id", BigInteger, primary_key=True, autoincrement=True),
Column("user_id", Text),
Column("endpoint", Text),
Column("timestamp", DateTime(timezone=True), nullable=False, server_default=func.now()),
Column("data", JSONB),
)
feedback_table = Table(
"feedback",
metadata,
Column("id", UUID(as_uuid=True), primary_key=True, server_default=func.gen_random_uuid()),
Column("conversation_id", UUID(as_uuid=True), nullable=False),
Column("user_id", Text, nullable=False),
Column("question_index", Integer, nullable=False),
Column("feedback_text", Text),
Column("timestamp", DateTime(timezone=True), nullable=False, server_default=func.now()),
)
stack_logs_table = Table(
"stack_logs",
metadata,
Column("id", BigInteger, primary_key=True, autoincrement=True),
Column("activity_id", Text, nullable=False),
Column("endpoint", Text),
Column("level", Text),
Column("user_id", Text),
Column("api_key", Text),
Column("query", Text),
Column("stacks", JSONB, nullable=False, server_default="[]"),
Column("timestamp", DateTime(timezone=True), nullable=False, server_default=func.now()),
)
@@ -0,0 +1,57 @@
"""Repository for the ``feedback`` table.
The ``feedback_collection`` global is declared in ``base.py`` but currently
has zero direct call sites in the application code (all feedback writes go
through ``conversation_messages.feedback`` JSONB field on the conversations
collection). The table exists for when feedback is denormalized into its own
rows. This repository provides the append-only insert and basic reads
needed for that future.
"""
from __future__ import annotations
from typing import Optional
from sqlalchemy import Connection, text
from application.storage.db.base_repository import row_to_dict
class FeedbackRepository:
"""Postgres-backed replacement for Mongo ``feedback_collection``."""
def __init__(self, conn: Connection) -> None:
self._conn = conn
def create(
self,
conversation_id: str,
user_id: str,
question_index: int,
feedback_text: Optional[str] = None,
) -> dict:
result = self._conn.execute(
text(
"""
INSERT INTO feedback (conversation_id, user_id, question_index, feedback_text)
VALUES (CAST(:conversation_id AS uuid), :user_id, :question_index, :feedback_text)
RETURNING *
"""
),
{
"conversation_id": conversation_id,
"user_id": user_id,
"question_index": question_index,
"feedback_text": feedback_text,
},
)
return row_to_dict(result.fetchone())
def list_for_conversation(self, conversation_id: str) -> list[dict]:
result = self._conn.execute(
text(
"SELECT * FROM feedback WHERE conversation_id = CAST(:cid AS uuid) ORDER BY question_index"
),
{"cid": conversation_id},
)
return [row_to_dict(r) for r in result.fetchall()]
@@ -0,0 +1,96 @@
"""Repository for the ``prompts`` table.
Covers every operation the legacy Mongo code performs on
``prompts_collection``:
1. ``insert_one`` in prompts/routes.py (create)
2. ``find`` by user in prompts/routes.py (list)
3. ``find_one`` by id+user in prompts/routes.py (get single)
4. ``find_one`` by id only in stream_processor.py (get content for rendering)
5. ``update_one`` in prompts/routes.py (update name+content)
6. ``delete_one`` in prompts/routes.py (delete)
7. ``find_one`` + ``insert_one`` in seeder.py (upsert by user+name+content)
"""
from __future__ import annotations
from typing import Optional
from sqlalchemy import Connection, text
from application.storage.db.base_repository import row_to_dict
class PromptsRepository:
"""Postgres-backed replacement for Mongo ``prompts_collection``."""
def __init__(self, conn: Connection) -> None:
self._conn = conn
def create(self, user_id: str, name: str, content: str) -> dict:
result = self._conn.execute(
text(
"""
INSERT INTO prompts (user_id, name, content)
VALUES (:user_id, :name, :content)
RETURNING *
"""
),
{"user_id": user_id, "name": name, "content": content},
)
return row_to_dict(result.fetchone())
def get(self, prompt_id: str, user_id: Optional[str] = None) -> Optional[dict]:
if user_id is not None:
result = self._conn.execute(
text("SELECT * FROM prompts WHERE id = CAST(:id AS uuid) AND user_id = :user_id"),
{"id": prompt_id, "user_id": user_id},
)
else:
result = self._conn.execute(
text("SELECT * FROM prompts WHERE id = CAST(:id AS uuid)"),
{"id": prompt_id},
)
row = result.fetchone()
return row_to_dict(row) if row is not None else None
def list_for_user(self, user_id: str) -> list[dict]:
result = self._conn.execute(
text("SELECT * FROM prompts WHERE user_id = :user_id ORDER BY created_at"),
{"user_id": user_id},
)
return [row_to_dict(r) for r in result.fetchall()]
def update(self, prompt_id: str, user_id: str, name: str, content: str) -> None:
self._conn.execute(
text(
"""
UPDATE prompts
SET name = :name, content = :content, updated_at = now()
WHERE id = CAST(:id AS uuid) AND user_id = :user_id
"""
),
{"id": prompt_id, "user_id": user_id, "name": name, "content": content},
)
def delete(self, prompt_id: str, user_id: str) -> None:
self._conn.execute(
text("DELETE FROM prompts WHERE id = CAST(:id AS uuid) AND user_id = :user_id"),
{"id": prompt_id, "user_id": user_id},
)
def find_or_create(self, user_id: str, name: str, content: str) -> dict:
"""Return existing prompt matching (user, name, content), or create one.
Used by the seeder to avoid duplicating template prompts.
"""
result = self._conn.execute(
text(
"SELECT * FROM prompts WHERE user_id = :user_id AND name = :name AND content = :content"
),
{"user_id": user_id, "name": name, "content": content},
)
row = result.fetchone()
if row is not None:
return row_to_dict(row)
return self.create(user_id, name, content)
@@ -0,0 +1,58 @@
"""Repository for the ``stack_logs`` table.
Covers the single operation the legacy Mongo code performs:
1. ``insert_one`` in logging.py ``_log_to_mongodb`` — append-only debug/error
activity log. The Mongo collection is ``stack_logs``; the Mongo variable
inside ``_log_to_mongodb`` is misleadingly named ``user_logs_collection``.
"""
from __future__ import annotations
import json
from datetime import datetime
from typing import Optional
from sqlalchemy import Connection, text
class StackLogsRepository:
"""Postgres-backed replacement for Mongo ``stack_logs`` collection."""
def __init__(self, conn: Connection) -> None:
self._conn = conn
def insert(
self,
*,
activity_id: str,
endpoint: Optional[str] = None,
level: Optional[str] = None,
user_id: Optional[str] = None,
api_key: Optional[str] = None,
query: Optional[str] = None,
stacks: Optional[list] = None,
timestamp: Optional[datetime] = None,
) -> None:
self._conn.execute(
text(
"""
INSERT INTO stack_logs (activity_id, endpoint, level, user_id, api_key, query, stacks, timestamp)
VALUES (
:activity_id, :endpoint, :level, :user_id, :api_key, :query,
CAST(:stacks AS jsonb),
COALESCE(:timestamp, now())
)
"""
),
{
"activity_id": activity_id,
"endpoint": endpoint,
"level": level,
"user_id": user_id,
"api_key": api_key,
"query": query,
"stacks": json.dumps(stacks or []),
"timestamp": timestamp,
},
)
@@ -0,0 +1,104 @@
"""Repository for the ``token_usage`` table.
Covers every operation the legacy Mongo code performs on
``token_usage_collection`` / ``usage_collection``:
1. ``insert_one`` in usage.py (record per-call token counts)
2. ``aggregate`` in analytics/routes.py (time-bucketed totals)
3. ``aggregate`` in answer/routes/base.py (24h sum for rate limiting)
4. ``count_documents`` in answer/routes/base.py (24h request count)
"""
from __future__ import annotations
from datetime import datetime
from typing import Optional
from sqlalchemy import Connection, text
class TokenUsageRepository:
"""Postgres-backed replacement for Mongo ``token_usage_collection``."""
def __init__(self, conn: Connection) -> None:
self._conn = conn
def insert(
self,
*,
user_id: Optional[str] = None,
api_key: Optional[str] = None,
agent_id: Optional[str] = None,
prompt_tokens: int = 0,
generated_tokens: int = 0,
timestamp: Optional[datetime] = None,
) -> None:
self._conn.execute(
text(
"""
INSERT INTO token_usage (user_id, api_key, agent_id, prompt_tokens, generated_tokens, timestamp)
VALUES (
:user_id, :api_key,
CAST(:agent_id AS uuid),
:prompt_tokens, :generated_tokens,
COALESCE(:timestamp, now())
)
"""
),
{
"user_id": user_id,
"api_key": api_key,
"agent_id": agent_id,
"prompt_tokens": prompt_tokens,
"generated_tokens": generated_tokens,
"timestamp": timestamp,
},
)
def sum_tokens_in_range(
self,
*,
start: datetime,
end: datetime,
user_id: Optional[str] = None,
api_key: Optional[str] = None,
) -> int:
"""Total (prompt + generated) tokens in the given time range."""
clauses = ["timestamp >= :start", "timestamp <= :end"]
params: dict = {"start": start, "end": end}
if user_id is not None:
clauses.append("user_id = :user_id")
params["user_id"] = user_id
if api_key is not None:
clauses.append("api_key = :api_key")
params["api_key"] = api_key
where = " AND ".join(clauses)
result = self._conn.execute(
text(f"SELECT COALESCE(SUM(prompt_tokens + generated_tokens), 0) FROM token_usage WHERE {where}"),
params,
)
return result.scalar()
def count_in_range(
self,
*,
start: datetime,
end: datetime,
user_id: Optional[str] = None,
api_key: Optional[str] = None,
) -> int:
"""Count of token_usage rows in the given time range (for request limiting)."""
clauses = ["timestamp >= :start", "timestamp <= :end"]
params: dict = {"start": start, "end": end}
if user_id is not None:
clauses.append("user_id = :user_id")
params["user_id"] = user_id
if api_key is not None:
clauses.append("api_key = :api_key")
params["api_key"] = api_key
where = " AND ".join(clauses)
result = self._conn.execute(
text(f"SELECT COUNT(*) FROM token_usage WHERE {where}"),
params,
)
return result.scalar()
@@ -0,0 +1,84 @@
"""Repository for the ``user_logs`` table.
Covers every operation the legacy Mongo code performs on
``user_logs_collection``:
1. ``insert_one`` in logging.py (per-request activity log via
``_log_to_mongodb`` — note: the *Mongo* variable is confusingly named
``user_logs_collection`` but points at the ``user_logs`` Mongo
collection, not ``stack_logs``)
2. ``insert_one`` in answer/routes/base.py (per-stream log entry)
3. ``find`` with sort/skip/limit in analytics/routes.py (paginated log list)
"""
from __future__ import annotations
import json
from datetime import datetime
from typing import Optional
from sqlalchemy import Connection, text
from application.storage.db.base_repository import row_to_dict
class UserLogsRepository:
"""Postgres-backed replacement for Mongo ``user_logs_collection``."""
def __init__(self, conn: Connection) -> None:
self._conn = conn
def insert(
self,
*,
user_id: Optional[str] = None,
endpoint: Optional[str] = None,
data: Optional[dict] = None,
timestamp: Optional[datetime] = None,
) -> None:
self._conn.execute(
text(
"""
INSERT INTO user_logs (user_id, endpoint, data, timestamp)
VALUES (:user_id, :endpoint, CAST(:data AS jsonb), COALESCE(:timestamp, now()))
"""
),
{
"user_id": user_id,
"endpoint": endpoint,
"data": json.dumps(data) if data is not None else None,
"timestamp": timestamp,
},
)
def list_paginated(
self,
*,
user_id: Optional[str] = None,
api_key: Optional[str] = None,
page: int = 1,
page_size: int = 10,
) -> tuple[list[dict], bool]:
"""Return ``(rows, has_more)`` for the requested page.
Mirrors the Mongo ``find(query).sort().skip().limit(page_size+1)``
pattern used in analytics/routes.py.
"""
clauses: list[str] = []
params: dict = {"limit": page_size + 1, "offset": (page - 1) * page_size}
if user_id is not None:
clauses.append("user_id = :user_id")
params["user_id"] = user_id
if api_key is not None:
clauses.append("data->>'api_key' = :api_key")
params["api_key"] = api_key
where = ("WHERE " + " AND ".join(clauses)) if clauses else ""
result = self._conn.execute(
text(
f"SELECT * FROM user_logs {where} ORDER BY timestamp DESC LIMIT :limit OFFSET :offset"
),
params,
)
rows = [row_to_dict(r) for r in result.fetchall()]
has_more = len(rows) > page_size
return rows[:page_size], has_more
@@ -0,0 +1,99 @@
"""Repository for the ``user_tools`` table.
Covers every operation the legacy Mongo code performs on
``user_tools_collection``:
1. ``find`` by user in tools/routes.py and base.py (list all / active)
2. ``find_one`` by id in tools/routes.py and sharing.py (get single)
3. ``insert_one`` in tools/routes.py and mcp.py (create)
4. ``update_one`` in tools/routes.py and mcp.py (update fields)
5. ``delete_one`` in tools/routes.py (delete)
6. ``find`` by user+status in stream_processor.py and tool_executor.py (active tools)
7. ``find_one`` by user+name in mcp.py (upsert check)
"""
from __future__ import annotations
import json
from typing import Optional
from sqlalchemy import Connection, text
from application.storage.db.base_repository import row_to_dict
class UserToolsRepository:
"""Postgres-backed replacement for Mongo ``user_tools_collection``."""
def __init__(self, conn: Connection) -> None:
self._conn = conn
def create(self, user_id: str, name: str, *, config: Optional[dict] = None,
custom_name: Optional[str] = None, display_name: Optional[str] = None,
extra: Optional[dict] = None) -> dict:
"""Insert a new tool row. ``extra`` is merged into the config JSONB."""
cfg = config or {}
if extra:
cfg.update(extra)
result = self._conn.execute(
text(
"""
INSERT INTO user_tools (user_id, name, custom_name, display_name, config)
VALUES (:user_id, :name, :custom_name, :display_name, CAST(:config AS jsonb))
RETURNING *
"""
),
{
"user_id": user_id,
"name": name,
"custom_name": custom_name,
"display_name": display_name,
"config": json.dumps(cfg),
},
)
return row_to_dict(result.fetchone())
def get(self, tool_id: str) -> Optional[dict]:
result = self._conn.execute(
text("SELECT * FROM user_tools WHERE id = CAST(:id AS uuid)"),
{"id": tool_id},
)
row = result.fetchone()
return row_to_dict(row) if row is not None else None
def list_for_user(self, user_id: str) -> list[dict]:
result = self._conn.execute(
text("SELECT * FROM user_tools WHERE user_id = :user_id ORDER BY created_at"),
{"user_id": user_id},
)
return [row_to_dict(r) for r in result.fetchall()]
def update(self, tool_id: str, user_id: str, fields: dict) -> None:
"""Update arbitrary fields on a tool row.
``fields`` maps column names to new values. Only ``name``,
``custom_name``, ``display_name``, and ``config`` are allowed.
"""
allowed = {"name", "custom_name", "display_name", "config"}
filtered = {k: v for k, v in fields.items() if k in allowed}
if not filtered:
return
set_clauses = []
params: dict = {"id": tool_id, "user_id": user_id}
for col, val in filtered.items():
if col == "config":
set_clauses.append(f"{col} = CAST(:val_{col} AS jsonb)")
params[f"val_{col}"] = json.dumps(val) if isinstance(val, dict) else val
else:
set_clauses.append(f"{col} = :val_{col}")
params[f"val_{col}"] = val
set_clauses.append("updated_at = now()")
sql = f"UPDATE user_tools SET {', '.join(set_clauses)} WHERE id = CAST(:id AS uuid) AND user_id = :user_id"
self._conn.execute(text(sql), params)
def delete(self, tool_id: str, user_id: str) -> bool:
result = self._conn.execute(
text("DELETE FROM user_tools WHERE id = CAST(:id AS uuid) AND user_id = :user_id"),
{"id": tool_id, "user_id": user_id},
)
return result.rowcount > 0
+14
View File
@@ -110,6 +110,20 @@ def update_token_usage(decoded_token, user_api_key, token_usage, agent_id=None):
usage_data["agent_id"] = normalized_agent_id
usage_collection.insert_one(usage_data)
from application.storage.db.dual_write import dual_write
from application.storage.db.repositories.token_usage import TokenUsageRepository
dual_write(
TokenUsageRepository,
lambda repo, d=usage_data: repo.insert(
user_id=d.get("user_id"),
api_key=d.get("api_key"),
agent_id=d.get("agent_id"),
prompt_tokens=d["prompt_tokens"],
generated_tokens=d["generated_tokens"],
),
)
def gen_token_usage(func):
def wrapper(self, model, messages, stream, tools, **kwargs):
+3 -1
View File
@@ -247,7 +247,7 @@ def extract_zip_recursive(zip_path, extract_to, current_depth=0, max_depth=5):
def download_file(url, params, dest_path):
try:
response = requests.get(url, params=params)
response = requests.get(url, params=params, timeout=100)
response.raise_for_status()
with open(dest_path, "wb") as f:
f.write(response.content)
@@ -284,12 +284,14 @@ def upload_index(full_path, file_data):
files=files,
data=file_data,
headers=headers,
timeout=100,
)
else:
response = requests.post(
urljoin(settings.API_URL, "/api/upload_index"),
data=file_data,
headers=headers,
timeout=100,
)
response.raise_for_status()
except (requests.RequestException, FileNotFoundError) as e:
+3 -3
View File
@@ -54,8 +54,8 @@ flowchart LR
* **Technology:** Supports multiple LLM APIs and local engines.
* **Responsibility:** This layer provides an abstraction for interacting with Large Language Models (LLMs).
* **Key Features:**
* Supports LLMs from OpenAI, Google, Anthropic, Groq, HuggingFace Inference API, Azure OpenAI, also compatable with local models like Ollama, LLaMa.cpp, Text Generation Inference (TGI), SGLang, vLLM, Aphrodite, FriendliAI, and LMDeploy.
* Manages API key handling and request formatting and Tool fromatting.
* Supports LLMs from OpenAI, Google, Anthropic, Groq, HuggingFace Inference API, Azure OpenAI, also compatible with local models like Ollama, LLaMa.cpp, Text Generation Inference (TGI), SGLang, vLLM, Aphrodite, FriendliAI, and LMDeploy.
* Manages API key handling and request formatting and Tool formatting.
* Offers caching mechanisms to improve response times and reduce API usage.
* Handles streaming responses for a more interactive user experience.
@@ -120,7 +120,7 @@ sequenceDiagram
## Deployment Architecture
DocsGPT is designed to be deployed using Docker and Kubernetes, here is a qucik overview of a simple k8s deployment.
DocsGPT is designed to be deployed using Docker and Kubernetes, here is a quick overview of a simple k8s deployment.
```mermaid
graph LR
+361 -2
View File
@@ -9,8 +9,8 @@ two-step change in this file:
2. Add a single entry to :data:`BACKFILLERS`.
There are intentionally no per-collection CLI flags or environment
variables — ``USE_POSTGRES`` / ``READ_POSTGRES`` in ``.env`` are the
only knobs operators need. This script discovers what's available from
variables — ``USE_POSTGRES`` in ``.env`` is the only knob operators
need during the migration window. This script discovers what's available from
the :data:`BACKFILLERS` registry and runs whichever tables were asked for.
Usage::
@@ -118,6 +118,359 @@ def _backfill_users(
return {"seen": seen, "written": written, "skipped_no_user_id": skipped}
def _backfill_prompts(
*,
conn: Connection,
mongo_db: Any,
batch_size: int,
dry_run: bool,
) -> dict:
"""Sync the ``prompts`` table from Mongo ``prompts`` collection."""
upsert_sql = text(
"""
INSERT INTO prompts (user_id, name, content)
VALUES (:user_id, :name, :content)
ON CONFLICT DO NOTHING
"""
)
cursor = (
mongo_db["prompts"]
.find({}, no_cursor_timeout=True)
.batch_size(batch_size)
)
seen = 0
written = 0
skipped = 0
batch: list[dict] = []
try:
for doc in cursor:
seen += 1
user_id = doc.get("user")
if not user_id:
skipped += 1
continue
batch.append({
"user_id": user_id,
"name": doc.get("name", ""),
"content": doc.get("content", ""),
})
if len(batch) >= batch_size:
if not dry_run:
conn.execute(upsert_sql, batch)
written += len(batch)
batch.clear()
if batch:
if not dry_run:
conn.execute(upsert_sql, batch)
written += len(batch)
finally:
cursor.close()
return {"seen": seen, "written": written, "skipped_no_user": skipped}
def _backfill_user_tools(
*,
conn: Connection,
mongo_db: Any,
batch_size: int,
dry_run: bool,
) -> dict:
"""Sync the ``user_tools`` table from Mongo ``user_tools`` collection."""
insert_sql = text(
"""
INSERT INTO user_tools (user_id, name, custom_name, display_name, config)
VALUES (:user_id, :name, :custom_name, :display_name, CAST(:config AS jsonb))
ON CONFLICT DO NOTHING
"""
)
cursor = (
mongo_db["user_tools"]
.find({}, no_cursor_timeout=True)
.batch_size(batch_size)
)
seen = 0
written = 0
skipped = 0
batch: list[dict] = []
try:
for doc in cursor:
seen += 1
user_id = doc.get("user")
if not user_id:
skipped += 1
continue
batch.append({
"user_id": user_id,
"name": doc.get("name", ""),
"custom_name": doc.get("customName"),
"display_name": doc.get("displayName"),
"config": json.dumps(doc.get("config") or {}),
})
if len(batch) >= batch_size:
if not dry_run:
conn.execute(insert_sql, batch)
written += len(batch)
batch.clear()
if batch:
if not dry_run:
conn.execute(insert_sql, batch)
written += len(batch)
finally:
cursor.close()
return {"seen": seen, "written": written, "skipped_no_user": skipped}
def _backfill_feedback(
*,
conn: Connection,
mongo_db: Any,
batch_size: int,
dry_run: bool,
) -> dict:
"""Sync the ``feedback`` table from Mongo ``feedback`` collection.
feedback.conversation_id is stored as a string UUID. Rows whose
conversation_id cannot be cast to UUID are skipped.
"""
insert_sql = text(
"""
INSERT INTO feedback (conversation_id, user_id, question_index, feedback_text, timestamp)
VALUES (CAST(:conversation_id AS uuid), :user_id, :question_index, :feedback_text, :timestamp)
ON CONFLICT DO NOTHING
"""
)
cursor = (
mongo_db["feedback"]
.find({}, no_cursor_timeout=True)
.batch_size(batch_size)
)
seen = 0
written = 0
skipped = 0
batch: list[dict] = []
try:
for doc in cursor:
seen += 1
user_id = doc.get("user")
conv_id = doc.get("conversation_id")
if not user_id or not conv_id:
skipped += 1
continue
batch.append({
"conversation_id": str(conv_id),
"user_id": user_id,
"question_index": doc.get("question_index", 0),
"feedback_text": doc.get("feedback_text"),
"timestamp": doc.get("timestamp"),
})
if len(batch) >= batch_size:
if not dry_run:
conn.execute(insert_sql, batch)
written += len(batch)
batch.clear()
if batch:
if not dry_run:
conn.execute(insert_sql, batch)
written += len(batch)
finally:
cursor.close()
return {"seen": seen, "written": written, "skipped": skipped}
def _backfill_stack_logs(
*,
conn: Connection,
mongo_db: Any,
batch_size: int,
dry_run: bool,
) -> dict:
"""Sync the ``stack_logs`` table from Mongo ``stack_logs`` collection."""
insert_sql = text(
"""
INSERT INTO stack_logs (activity_id, endpoint, level, user_id, api_key, query, stacks, timestamp)
VALUES (:activity_id, :endpoint, :level, :user_id, :api_key, :query, CAST(:stacks AS jsonb), :timestamp)
"""
)
cursor = (
mongo_db["stack_logs"]
.find({}, no_cursor_timeout=True)
.batch_size(batch_size)
)
seen = 0
written = 0
skipped = 0
batch: list[dict] = []
try:
for doc in cursor:
seen += 1
activity_id = doc.get("id")
if not activity_id:
skipped += 1
continue
batch.append({
"activity_id": str(activity_id),
"endpoint": doc.get("endpoint"),
"level": doc.get("level"),
"user_id": doc.get("user"),
"api_key": doc.get("api_key"),
"query": doc.get("query"),
"stacks": json.dumps(doc.get("stacks") or []),
"timestamp": doc.get("timestamp"),
})
if len(batch) >= batch_size:
if not dry_run:
conn.execute(insert_sql, batch)
written += len(batch)
batch.clear()
if batch:
if not dry_run:
conn.execute(insert_sql, batch)
written += len(batch)
finally:
cursor.close()
return {"seen": seen, "written": written, "skipped_no_id": skipped}
def _backfill_user_logs(
*,
conn: Connection,
mongo_db: Any,
batch_size: int,
dry_run: bool,
) -> dict:
"""Sync the ``user_logs`` table from Mongo ``user_logs`` collection."""
insert_sql = text(
"""
INSERT INTO user_logs (user_id, endpoint, data, timestamp)
VALUES (:user_id, :endpoint, CAST(:data AS jsonb), :timestamp)
"""
)
cursor = (
mongo_db["user_logs"]
.find({}, no_cursor_timeout=True)
.batch_size(batch_size)
)
seen = 0
written = 0
batch: list[dict] = []
try:
for doc in cursor:
seen += 1
# Build a JSONB payload from the full doc (minus Mongo internals).
data_payload = {k: v for k, v in doc.items() if k != "_id"}
# Stringify ObjectId values inside the payload.
for k, v in data_payload.items():
if hasattr(v, "__str__") and type(v).__name__ == "ObjectId":
data_payload[k] = str(v)
batch.append({
"user_id": doc.get("user"),
"endpoint": doc.get("action") or doc.get("endpoint"),
"data": json.dumps(data_payload, default=str),
"timestamp": doc.get("timestamp"),
})
if len(batch) >= batch_size:
if not dry_run:
conn.execute(insert_sql, batch)
written += len(batch)
batch.clear()
if batch:
if not dry_run:
conn.execute(insert_sql, batch)
written += len(batch)
finally:
cursor.close()
return {"seen": seen, "written": written}
def _backfill_token_usage(
*,
conn: Connection,
mongo_db: Any,
batch_size: int,
dry_run: bool,
) -> dict:
"""Sync the ``token_usage`` table from Mongo ``token_usage`` collection."""
insert_sql = text(
"""
INSERT INTO token_usage (user_id, api_key, agent_id, prompt_tokens, generated_tokens, timestamp)
VALUES (
:user_id, :api_key,
CAST(:agent_id AS uuid),
:prompt_tokens, :generated_tokens, :timestamp
)
"""
)
cursor = (
mongo_db["token_usage"]
.find({}, no_cursor_timeout=True)
.batch_size(batch_size)
)
seen = 0
written = 0
batch: list[dict] = []
try:
for doc in cursor:
seen += 1
agent_id = doc.get("agent_id")
# agent_id may be an ObjectId string or None — only pass if
# it looks like a valid UUID (from dual-write) or skip it.
agent_id_str = None
if agent_id:
s = str(agent_id)
if len(s) == 36 and "-" in s:
agent_id_str = s
batch.append({
"user_id": doc.get("user_id"),
"api_key": doc.get("api_key"),
"agent_id": agent_id_str,
"prompt_tokens": doc.get("prompt_tokens", 0),
"generated_tokens": doc.get("generated_tokens", 0),
"timestamp": doc.get("timestamp"),
})
if len(batch) >= batch_size:
if not dry_run:
conn.execute(insert_sql, batch)
written += len(batch)
batch.clear()
if batch:
if not dry_run:
conn.execute(insert_sql, batch)
written += len(batch)
finally:
cursor.close()
return {"seen": seen, "written": written}
# ---------------------------------------------------------------------------
# Registry
# ---------------------------------------------------------------------------
@@ -131,6 +484,12 @@ BackfillFn = Callable[..., dict]
# backfill doesn't hit FK errors.
BACKFILLERS: dict[str, BackfillFn] = {
"users": _backfill_users,
"prompts": _backfill_prompts,
"user_tools": _backfill_user_tools,
"feedback": _backfill_feedback,
"stack_logs": _backfill_stack_logs,
"user_logs": _backfill_user_logs,
"token_usage": _backfill_token_usage,
}
View File
Whitespace-only changes.
+66
View File
@@ -0,0 +1,66 @@
"""Fixtures for repository tests against a real Postgres instance.
These tests hit the local dev Postgres (the DBngin instance on this machine,
or CI's service container). Each test runs inside a transaction that is
rolled back at the end, so tests never leak state into each other and the
database stays clean without needing per-test CREATE/DROP overhead.
Required env:
POSTGRES_URI — e.g. postgresql+psycopg://docsgpt:docsgpt@localhost:5432/docsgpt
Tests are skipped automatically when POSTGRES_URI is unset so that
contributors without a local Postgres can still run the rest of the suite.
"""
from __future__ import annotations
import subprocess
import sys
from pathlib import Path
import pytest
from sqlalchemy import create_engine, text
from application.core.settings import settings
def _run_alembic_upgrade(engine):
"""Run ``alembic upgrade head`` to ensure the full schema is present.
Falls back to inline DDL for CI environments where alembic is not
on PATH (shouldn't happen, but defence in depth).
"""
alembic_ini = Path(__file__).resolve().parents[3] / "application" / "alembic.ini"
try:
subprocess.check_call(
[sys.executable, "-m", "alembic", "-c", str(alembic_ini), "upgrade", "head"],
timeout=30,
)
except Exception:
# Alembic failed — tables likely already exist from a prior run.
pass
@pytest.fixture(scope="session")
def pg_engine():
"""Session-scoped engine pointing at the test Postgres."""
if not settings.POSTGRES_URI:
pytest.skip("POSTGRES_URI not set")
engine = create_engine(settings.POSTGRES_URI)
_run_alembic_upgrade(engine)
yield engine
engine.dispose()
@pytest.fixture()
def pg_conn(pg_engine):
"""Per-test connection wrapped in a transaction that always rolls back.
Repositories receive this connection and operate normally. At teardown
the outer transaction is rolled back so no data persists between tests.
"""
conn = pg_engine.connect()
txn = conn.begin()
yield conn
txn.rollback()
conn.close()
Whitespace-only changes.
@@ -0,0 +1,79 @@
"""Tests for FeedbackRepository against a real Postgres instance."""
from __future__ import annotations
import uuid
import pytest
from sqlalchemy import text
from application.storage.db.repositories.feedback import FeedbackRepository
pytestmark = pytest.mark.skipif(
not __import__("application.core.settings", fromlist=["settings"]).settings.POSTGRES_URI,
reason="POSTGRES_URI not configured",
)
def _repo(conn) -> FeedbackRepository:
return FeedbackRepository(conn)
def _make_conversation_id(pg_conn) -> str:
"""Insert a minimal conversations row and return its id as string.
feedback has a FK to conversations (added in Tier 3 migration), but
the FK constraint may not exist yet during early phases. We create a
row anyway to keep tests realistic.
"""
cid = str(uuid.uuid4())
# Only insert if the conversations table exists; otherwise use a random UUID.
row = pg_conn.execute(
text(
"SELECT 1 FROM information_schema.tables "
"WHERE table_schema='public' AND table_name='conversations'"
)
).scalar()
if row:
pg_conn.execute(
text("INSERT INTO conversations (id, user_id) VALUES (CAST(:id AS uuid), 'test')"),
{"id": cid},
)
return cid
class TestCreate:
def test_creates_feedback(self, pg_conn):
repo = _repo(pg_conn)
cid = _make_conversation_id(pg_conn)
doc = repo.create(cid, "user-1", 0, "great answer")
assert doc["conversation_id"] is not None
assert doc["user_id"] == "user-1"
assert doc["question_index"] == 0
assert doc["feedback_text"] == "great answer"
def test_allows_null_feedback_text(self, pg_conn):
repo = _repo(pg_conn)
cid = _make_conversation_id(pg_conn)
doc = repo.create(cid, "user-1", 1)
assert doc["feedback_text"] is None
class TestListForConversation:
def test_lists_feedback_for_conversation(self, pg_conn):
repo = _repo(pg_conn)
cid = _make_conversation_id(pg_conn)
repo.create(cid, "user-1", 0, "good")
repo.create(cid, "user-1", 1, "bad")
results = repo.list_for_conversation(cid)
assert len(results) == 2
assert results[0]["question_index"] == 0
assert results[1]["question_index"] == 1
def test_does_not_mix_conversations(self, pg_conn):
repo = _repo(pg_conn)
cid1 = _make_conversation_id(pg_conn)
cid2 = _make_conversation_id(pg_conn)
repo.create(cid1, "user-1", 0, "a")
repo.create(cid2, "user-1", 0, "b")
assert len(repo.list_for_conversation(cid1)) == 1
@@ -0,0 +1,115 @@
"""Tests for PromptsRepository against a real Postgres instance."""
from __future__ import annotations
import pytest
from application.storage.db.repositories.prompts import PromptsRepository
pytestmark = pytest.mark.skipif(
not __import__("application.core.settings", fromlist=["settings"]).settings.POSTGRES_URI,
reason="POSTGRES_URI not configured",
)
def _repo(conn) -> PromptsRepository:
return PromptsRepository(conn)
class TestCreate:
def test_creates_prompt(self, pg_conn):
repo = _repo(pg_conn)
doc = repo.create("user-1", "greeting", "Hello {{name}}")
assert doc["user_id"] == "user-1"
assert doc["name"] == "greeting"
assert doc["content"] == "Hello {{name}}"
assert doc["id"] is not None
def test_create_returns_id_and_underscore_id(self, pg_conn):
repo = _repo(pg_conn)
doc = repo.create("user-1", "p", "c")
assert doc["_id"] == doc["id"]
class TestGet:
def test_get_by_id_and_user(self, pg_conn):
repo = _repo(pg_conn)
created = repo.create("user-1", "p", "c")
fetched = repo.get(created["id"], user_id="user-1")
assert fetched["id"] == created["id"]
def test_get_by_id_only(self, pg_conn):
repo = _repo(pg_conn)
created = repo.create("user-1", "p", "c")
fetched = repo.get(created["id"])
assert fetched is not None
def test_get_wrong_user_returns_none(self, pg_conn):
repo = _repo(pg_conn)
created = repo.create("user-1", "p", "c")
assert repo.get(created["id"], user_id="user-other") is None
def test_get_nonexistent_returns_none(self, pg_conn):
repo = _repo(pg_conn)
assert repo.get("00000000-0000-0000-0000-000000000000") is None
class TestListForUser:
def test_lists_only_own_prompts(self, pg_conn):
repo = _repo(pg_conn)
repo.create("alice", "a1", "c1")
repo.create("alice", "a2", "c2")
repo.create("bob", "b1", "c3")
results = repo.list_for_user("alice")
assert len(results) == 2
assert all(r["user_id"] == "alice" for r in results)
class TestUpdate:
def test_updates_name_and_content(self, pg_conn):
repo = _repo(pg_conn)
created = repo.create("user-1", "old", "old-content")
repo.update(created["id"], "user-1", "new", "new-content")
fetched = repo.get(created["id"])
assert fetched["name"] == "new"
assert fetched["content"] == "new-content"
def test_update_wrong_user_is_noop(self, pg_conn):
repo = _repo(pg_conn)
created = repo.create("user-1", "old", "old-content")
repo.update(created["id"], "user-other", "new", "new-content")
fetched = repo.get(created["id"])
assert fetched["name"] == "old"
class TestDelete:
def test_deletes_prompt(self, pg_conn):
repo = _repo(pg_conn)
created = repo.create("user-1", "p", "c")
repo.delete(created["id"], "user-1")
assert repo.get(created["id"]) is None
def test_delete_wrong_user_is_noop(self, pg_conn):
repo = _repo(pg_conn)
created = repo.create("user-1", "p", "c")
repo.delete(created["id"], "user-other")
assert repo.get(created["id"]) is not None
class TestFindOrCreate:
def test_creates_when_missing(self, pg_conn):
repo = _repo(pg_conn)
doc = repo.find_or_create("sys", "template", "content")
assert doc["id"] is not None
def test_returns_existing_on_match(self, pg_conn):
repo = _repo(pg_conn)
first = repo.find_or_create("sys", "template", "content")
second = repo.find_or_create("sys", "template", "content")
assert first["id"] == second["id"]
def test_different_content_creates_new(self, pg_conn):
repo = _repo(pg_conn)
first = repo.find_or_create("sys", "template", "v1")
second = repo.find_or_create("sys", "template", "v2")
assert first["id"] != second["id"]
@@ -0,0 +1,58 @@
"""Tests for StackLogsRepository against a real Postgres instance."""
from __future__ import annotations
import pytest
from sqlalchemy import text
from application.storage.db.repositories.stack_logs import StackLogsRepository
pytestmark = pytest.mark.skipif(
not __import__("application.core.settings", fromlist=["settings"]).settings.POSTGRES_URI,
reason="POSTGRES_URI not configured",
)
def _repo(conn) -> StackLogsRepository:
return StackLogsRepository(conn)
class TestInsert:
def test_inserts_log(self, pg_conn):
repo = _repo(pg_conn)
repo.insert(
activity_id="act-1",
endpoint="/api/answer",
level="info",
user_id="u1",
api_key="k1",
query="what is python?",
stacks=[{"component": "retriever", "data": {"docs": 3}}],
)
row = pg_conn.execute(
text("SELECT * FROM stack_logs WHERE activity_id = 'act-1'")
).fetchone()
assert row is not None
mapping = dict(row._mapping)
assert mapping["endpoint"] == "/api/answer"
assert mapping["level"] == "info"
assert mapping["user_id"] == "u1"
assert mapping["stacks"] == [{"component": "retriever", "data": {"docs": 3}}]
def test_inserts_with_empty_stacks(self, pg_conn):
repo = _repo(pg_conn)
repo.insert(activity_id="act-2", level="error")
row = pg_conn.execute(
text("SELECT stacks FROM stack_logs WHERE activity_id = 'act-2'")
).fetchone()
assert row is not None
assert dict(row._mapping)["stacks"] == []
def test_truncated_query_stored(self, pg_conn):
repo = _repo(pg_conn)
long_query = "x" * 20000
repo.insert(activity_id="act-3", query=long_query)
row = pg_conn.execute(
text("SELECT query FROM stack_logs WHERE activity_id = 'act-3'")
).fetchone()
assert len(dict(row._mapping)["query"]) == 20000
@@ -0,0 +1,90 @@
"""Tests for TokenUsageRepository against a real Postgres instance."""
from __future__ import annotations
from datetime import datetime, timedelta, timezone
import pytest
from application.storage.db.repositories.token_usage import TokenUsageRepository
pytestmark = pytest.mark.skipif(
not __import__("application.core.settings", fromlist=["settings"]).settings.POSTGRES_URI,
reason="POSTGRES_URI not configured",
)
def _repo(conn) -> TokenUsageRepository:
return TokenUsageRepository(conn)
def _now():
return datetime.now(timezone.utc)
class TestInsert:
def test_inserts_row(self, pg_conn):
repo = _repo(pg_conn)
repo.insert(user_id="u1", prompt_tokens=10, generated_tokens=5)
total = repo.sum_tokens_in_range(
start=_now() - timedelta(minutes=1), end=_now() + timedelta(minutes=1), user_id="u1"
)
assert total == 15
def test_insert_with_api_key(self, pg_conn):
repo = _repo(pg_conn)
repo.insert(api_key="key-1", prompt_tokens=20, generated_tokens=10)
total = repo.sum_tokens_in_range(
start=_now() - timedelta(minutes=1), end=_now() + timedelta(minutes=1), api_key="key-1"
)
assert total == 30
class TestSumTokensInRange:
def test_sums_correctly(self, pg_conn):
repo = _repo(pg_conn)
repo.insert(user_id="u1", prompt_tokens=10, generated_tokens=5)
repo.insert(user_id="u1", prompt_tokens=20, generated_tokens=10)
repo.insert(user_id="u2", prompt_tokens=100, generated_tokens=50)
total = repo.sum_tokens_in_range(
start=_now() - timedelta(minutes=1), end=_now() + timedelta(minutes=1), user_id="u1"
)
assert total == 45
def test_returns_zero_when_no_rows(self, pg_conn):
repo = _repo(pg_conn)
total = repo.sum_tokens_in_range(
start=_now() - timedelta(minutes=1), end=_now() + timedelta(minutes=1), user_id="nobody"
)
assert total == 0
def test_respects_time_range(self, pg_conn):
repo = _repo(pg_conn)
old = _now() - timedelta(hours=48)
repo.insert(user_id="u1", prompt_tokens=100, generated_tokens=0, timestamp=old)
repo.insert(user_id="u1", prompt_tokens=10, generated_tokens=0)
total = repo.sum_tokens_in_range(
start=_now() - timedelta(hours=1), end=_now() + timedelta(minutes=1), user_id="u1"
)
assert total == 10
class TestCountInRange:
def test_counts_rows(self, pg_conn):
repo = _repo(pg_conn)
repo.insert(user_id="u1", prompt_tokens=1, generated_tokens=1)
repo.insert(user_id="u1", prompt_tokens=1, generated_tokens=1)
repo.insert(user_id="u2", prompt_tokens=1, generated_tokens=1)
count = repo.count_in_range(
start=_now() - timedelta(minutes=1), end=_now() + timedelta(minutes=1), user_id="u1"
)
assert count == 2
def test_filters_by_api_key(self, pg_conn):
repo = _repo(pg_conn)
repo.insert(api_key="k1", prompt_tokens=1, generated_tokens=1)
repo.insert(api_key="k2", prompt_tokens=1, generated_tokens=1)
count = repo.count_in_range(
start=_now() - timedelta(minutes=1), end=_now() + timedelta(minutes=1), api_key="k1"
)
assert count == 1
@@ -0,0 +1,64 @@
"""Tests for UserLogsRepository against a real Postgres instance."""
from __future__ import annotations
import pytest
from application.storage.db.repositories.user_logs import UserLogsRepository
pytestmark = pytest.mark.skipif(
not __import__("application.core.settings", fromlist=["settings"]).settings.POSTGRES_URI,
reason="POSTGRES_URI not configured",
)
def _repo(conn) -> UserLogsRepository:
return UserLogsRepository(conn)
class TestInsert:
def test_inserts_log(self, pg_conn):
repo = _repo(pg_conn)
repo.insert(user_id="u1", endpoint="/api/answer", data={"question": "hi"})
rows, _ = repo.list_paginated(user_id="u1")
assert len(rows) == 1
assert rows[0]["data"]["question"] == "hi"
def test_allows_null_data(self, pg_conn):
repo = _repo(pg_conn)
repo.insert(user_id="u1")
rows, _ = repo.list_paginated(user_id="u1")
assert len(rows) == 1
assert rows[0]["data"] is None
class TestListPaginated:
def test_paginates(self, pg_conn):
repo = _repo(pg_conn)
for i in range(5):
repo.insert(user_id="u1", data={"i": i})
page1, has_more1 = repo.list_paginated(user_id="u1", page=1, page_size=3)
assert len(page1) == 3
assert has_more1 is True
page2, has_more2 = repo.list_paginated(user_id="u1", page=2, page_size=3)
assert len(page2) == 2
assert has_more2 is False
def test_filters_by_user(self, pg_conn):
repo = _repo(pg_conn)
repo.insert(user_id="alice", data={"x": 1})
repo.insert(user_id="bob", data={"x": 2})
rows, _ = repo.list_paginated(user_id="alice")
assert len(rows) == 1
assert rows[0]["user_id"] == "alice"
def test_ordered_by_timestamp_desc(self, pg_conn):
from datetime import datetime, timedelta, timezone
repo = _repo(pg_conn)
earlier = datetime.now(timezone.utc) - timedelta(minutes=5)
later = datetime.now(timezone.utc)
repo.insert(user_id="u1", data={"order": "first"}, timestamp=earlier)
repo.insert(user_id="u1", data={"order": "second"}, timestamp=later)
rows, _ = repo.list_paginated(user_id="u1")
assert rows[0]["data"]["order"] == "second"
@@ -0,0 +1,98 @@
"""Tests for UserToolsRepository against a real Postgres instance."""
from __future__ import annotations
import pytest
from application.storage.db.repositories.user_tools import UserToolsRepository
pytestmark = pytest.mark.skipif(
not __import__("application.core.settings", fromlist=["settings"]).settings.POSTGRES_URI,
reason="POSTGRES_URI not configured",
)
def _repo(conn) -> UserToolsRepository:
return UserToolsRepository(conn)
class TestCreate:
def test_creates_tool(self, pg_conn):
repo = _repo(pg_conn)
doc = repo.create("user-1", "my_tool", config={"key": "val"})
assert doc["user_id"] == "user-1"
assert doc["name"] == "my_tool"
assert doc["config"] == {"key": "val"}
assert doc["id"] is not None
def test_create_with_display_names(self, pg_conn):
repo = _repo(pg_conn)
doc = repo.create("user-1", "t", custom_name="Custom", display_name="Display")
assert doc["custom_name"] == "Custom"
assert doc["display_name"] == "Display"
class TestGet:
def test_get_existing(self, pg_conn):
repo = _repo(pg_conn)
created = repo.create("user-1", "t")
fetched = repo.get(created["id"])
assert fetched["id"] == created["id"]
def test_get_nonexistent(self, pg_conn):
repo = _repo(pg_conn)
assert repo.get("00000000-0000-0000-0000-000000000000") is None
class TestListForUser:
def test_lists_only_own_tools(self, pg_conn):
repo = _repo(pg_conn)
repo.create("alice", "t1")
repo.create("alice", "t2")
repo.create("bob", "t3")
results = repo.list_for_user("alice")
assert len(results) == 2
class TestUpdate:
def test_updates_name(self, pg_conn):
repo = _repo(pg_conn)
created = repo.create("user-1", "old_name")
repo.update(created["id"], "user-1", {"name": "new_name"})
fetched = repo.get(created["id"])
assert fetched["name"] == "new_name"
def test_updates_config(self, pg_conn):
repo = _repo(pg_conn)
created = repo.create("user-1", "t", config={"a": 1})
repo.update(created["id"], "user-1", {"config": {"a": 2, "b": 3}})
fetched = repo.get(created["id"])
assert fetched["config"] == {"a": 2, "b": 3}
def test_update_wrong_user_is_noop(self, pg_conn):
repo = _repo(pg_conn)
created = repo.create("user-1", "old")
repo.update(created["id"], "user-other", {"name": "new"})
fetched = repo.get(created["id"])
assert fetched["name"] == "old"
def test_ignores_disallowed_fields(self, pg_conn):
repo = _repo(pg_conn)
created = repo.create("user-1", "t")
repo.update(created["id"], "user-1", {"id": "00000000-0000-0000-0000-000000000000"})
fetched = repo.get(created["id"])
assert fetched["id"] == created["id"]
class TestDelete:
def test_deletes_tool(self, pg_conn):
repo = _repo(pg_conn)
created = repo.create("user-1", "t")
assert repo.delete(created["id"], "user-1") is True
assert repo.get(created["id"]) is None
def test_delete_wrong_user_returns_false(self, pg_conn):
repo = _repo(pg_conn)
created = repo.create("user-1", "t")
assert repo.delete(created["id"], "user-other") is False
assert repo.get(created["id"]) is not None
+222
View File
@@ -0,0 +1,222 @@
"""Tests for UsersRepository against a real Postgres instance.
Every test runs inside a rolled-back transaction (see ``pg_conn`` fixture
in the parent conftest) so no data leaks between tests.
"""
from __future__ import annotations
import pytest
from application.storage.db.repositories.users import UsersRepository
pytestmark = pytest.mark.skipif(
not __import__("application.core.settings", fromlist=["settings"]).settings.POSTGRES_URI,
reason="POSTGRES_URI not configured",
)
# ------------------------------------------------------------------
# Helpers
# ------------------------------------------------------------------
def _repo(conn) -> UsersRepository:
return UsersRepository(conn)
# ------------------------------------------------------------------
# upsert / get
# ------------------------------------------------------------------
class TestUpsert:
def test_creates_new_user_with_defaults(self, pg_conn):
repo = _repo(pg_conn)
doc = repo.upsert("user-new")
assert doc["user_id"] == "user-new"
assert doc["agent_preferences"] == {"pinned": [], "shared_with_me": []}
assert "id" in doc
assert doc["_id"] == doc["id"]
def test_upsert_is_idempotent(self, pg_conn):
repo = _repo(pg_conn)
first = repo.upsert("user-idem")
second = repo.upsert("user-idem")
assert first["id"] == second["id"]
assert first["agent_preferences"] == second["agent_preferences"]
def test_upsert_preserves_existing_preferences(self, pg_conn):
repo = _repo(pg_conn)
repo.upsert("user-prefs")
repo.add_pinned("user-prefs", "agent-1")
doc = repo.upsert("user-prefs")
assert "agent-1" in doc["agent_preferences"]["pinned"]
class TestGet:
def test_returns_none_for_missing_user(self, pg_conn):
repo = _repo(pg_conn)
assert repo.get("nonexistent") is None
def test_returns_user_after_upsert(self, pg_conn):
repo = _repo(pg_conn)
repo.upsert("user-get")
doc = repo.get("user-get")
assert doc is not None
assert doc["user_id"] == "user-get"
# ------------------------------------------------------------------
# pinned agents
# ------------------------------------------------------------------
class TestPinned:
def test_add_pinned(self, pg_conn):
repo = _repo(pg_conn)
repo.upsert("user-pin")
repo.add_pinned("user-pin", "a1")
doc = repo.get("user-pin")
assert doc["agent_preferences"]["pinned"] == ["a1"]
def test_add_pinned_is_idempotent(self, pg_conn):
repo = _repo(pg_conn)
repo.upsert("user-pin2")
repo.add_pinned("user-pin2", "a1")
repo.add_pinned("user-pin2", "a1")
doc = repo.get("user-pin2")
assert doc["agent_preferences"]["pinned"] == ["a1"]
def test_add_multiple_pinned(self, pg_conn):
repo = _repo(pg_conn)
repo.upsert("user-pin3")
repo.add_pinned("user-pin3", "a1")
repo.add_pinned("user-pin3", "a2")
doc = repo.get("user-pin3")
assert set(doc["agent_preferences"]["pinned"]) == {"a1", "a2"}
def test_remove_pinned(self, pg_conn):
repo = _repo(pg_conn)
repo.upsert("user-unpin")
repo.add_pinned("user-unpin", "a1")
repo.add_pinned("user-unpin", "a2")
repo.remove_pinned("user-unpin", "a1")
doc = repo.get("user-unpin")
assert doc["agent_preferences"]["pinned"] == ["a2"]
def test_remove_pinned_nonexistent_is_noop(self, pg_conn):
repo = _repo(pg_conn)
repo.upsert("user-unpin2")
repo.add_pinned("user-unpin2", "a1")
repo.remove_pinned("user-unpin2", "zzz")
doc = repo.get("user-unpin2")
assert doc["agent_preferences"]["pinned"] == ["a1"]
def test_remove_pinned_bulk(self, pg_conn):
repo = _repo(pg_conn)
repo.upsert("user-bulk")
repo.add_pinned("user-bulk", "a1")
repo.add_pinned("user-bulk", "a2")
repo.add_pinned("user-bulk", "a3")
repo.remove_pinned_bulk("user-bulk", ["a1", "a3"])
doc = repo.get("user-bulk")
assert doc["agent_preferences"]["pinned"] == ["a2"]
def test_remove_pinned_bulk_empty_list_is_noop(self, pg_conn):
repo = _repo(pg_conn)
repo.upsert("user-bulk2")
repo.add_pinned("user-bulk2", "a1")
repo.remove_pinned_bulk("user-bulk2", [])
doc = repo.get("user-bulk2")
assert doc["agent_preferences"]["pinned"] == ["a1"]
# ------------------------------------------------------------------
# shared_with_me
# ------------------------------------------------------------------
class TestShared:
def test_add_shared(self, pg_conn):
repo = _repo(pg_conn)
repo.upsert("user-share")
repo.add_shared("user-share", "s1")
doc = repo.get("user-share")
assert doc["agent_preferences"]["shared_with_me"] == ["s1"]
def test_add_shared_is_idempotent(self, pg_conn):
repo = _repo(pg_conn)
repo.upsert("user-share2")
repo.add_shared("user-share2", "s1")
repo.add_shared("user-share2", "s1")
doc = repo.get("user-share2")
assert doc["agent_preferences"]["shared_with_me"] == ["s1"]
def test_remove_shared_bulk(self, pg_conn):
repo = _repo(pg_conn)
repo.upsert("user-rshare")
repo.add_shared("user-rshare", "s1")
repo.add_shared("user-rshare", "s2")
repo.remove_shared_bulk("user-rshare", ["s1"])
doc = repo.get("user-rshare")
assert doc["agent_preferences"]["shared_with_me"] == ["s2"]
# ------------------------------------------------------------------
# remove_agent_from_all (cascade on agent delete)
# ------------------------------------------------------------------
class TestRemoveAgentFromAll:
def test_removes_from_both_arrays(self, pg_conn):
repo = _repo(pg_conn)
repo.upsert("user-cascade")
repo.add_pinned("user-cascade", "agent-x")
repo.add_shared("user-cascade", "agent-x")
repo.remove_agent_from_all("user-cascade", "agent-x")
doc = repo.get("user-cascade")
assert "agent-x" not in doc["agent_preferences"]["pinned"]
assert "agent-x" not in doc["agent_preferences"]["shared_with_me"]
def test_leaves_other_agents_untouched(self, pg_conn):
repo = _repo(pg_conn)
repo.upsert("user-cascade2")
repo.add_pinned("user-cascade2", "keep")
repo.add_pinned("user-cascade2", "remove")
repo.add_shared("user-cascade2", "keep")
repo.add_shared("user-cascade2", "remove")
repo.remove_agent_from_all("user-cascade2", "remove")
doc = repo.get("user-cascade2")
assert doc["agent_preferences"]["pinned"] == ["keep"]
assert doc["agent_preferences"]["shared_with_me"] == ["keep"]
def test_noop_when_agent_not_present(self, pg_conn):
repo = _repo(pg_conn)
repo.upsert("user-cascade3")
repo.add_pinned("user-cascade3", "a1")
repo.remove_agent_from_all("user-cascade3", "nonexistent")
doc = repo.get("user-cascade3")
assert doc["agent_preferences"]["pinned"] == ["a1"]
# ------------------------------------------------------------------
# tenant isolation
# ------------------------------------------------------------------
class TestTenantIsolation:
def test_get_cannot_see_other_users(self, pg_conn):
repo = _repo(pg_conn)
repo.upsert("alice")
repo.upsert("bob")
repo.add_pinned("alice", "a-private")
bob_doc = repo.get("bob")
assert "a-private" not in bob_doc["agent_preferences"]["pinned"]
def test_mutations_on_one_user_dont_affect_another(self, pg_conn):
repo = _repo(pg_conn)
repo.upsert("alice")
repo.upsert("bob")
repo.add_pinned("alice", "a1")
repo.add_shared("bob", "s1")
alice = repo.get("alice")
bob = repo.get("bob")
assert alice["agent_preferences"]["pinned"] == ["a1"]
assert alice["agent_preferences"]["shared_with_me"] == []
assert bob["agent_preferences"]["pinned"] == []
assert bob["agent_preferences"]["shared_with_me"] == ["s1"]