From 502819ae52637c96fc6fa376d5708e3e4f2b0e8d Mon Sep 17 00:00:00 2001 From: Alex Date: Sun, 12 Apr 2026 12:15:59 +0100 Subject: [PATCH 1/5] feat: pg migration, more tables --- .../config/vocabularies/DocsGPT/accept.txt | 123 +++--- application/agents/tools/brave.py | 4 +- application/agents/tools/cryptoprice.py | 2 +- application/agents/tools/ntfy.py | 2 +- application/agents/tools/telegram.py | 4 +- application/api/answer/routes/base.py | 12 + application/api/user/prompts/routes.py | 14 + application/api/user/tools/routes.py | 13 + application/core/settings.py | 8 +- application/logging.py | 16 + .../parser/connectors/share_point/loader.py | 21 +- application/parser/file/docs_parser.py | 2 +- application/parser/file/image_parser.py | 2 +- application/parser/remote/github_loader.py | 2 +- application/storage/db/engine.py | 41 +- application/storage/db/models.py | 73 ++++ .../storage/db/repositories/feedback.py | 57 +++ .../storage/db/repositories/prompts.py | 96 +++++ .../storage/db/repositories/stack_logs.py | 58 +++ .../storage/db/repositories/token_usage.py | 104 +++++ .../storage/db/repositories/user_logs.py | 84 ++++ .../storage/db/repositories/user_tools.py | 99 +++++ application/usage.py | 14 + application/worker.py | 4 +- docs/content/Guides/Architecture.mdx | 6 +- scripts/db/backfill.py | 363 +++++++++++++++++- tests/storage/db/__init__.py | 0 tests/storage/db/conftest.py | 66 ++++ tests/storage/db/repositories/__init__.py | 0 .../storage/db/repositories/test_feedback.py | 79 ++++ tests/storage/db/repositories/test_prompts.py | 115 ++++++ .../db/repositories/test_stack_logs.py | 58 +++ .../db/repositories/test_token_usage.py | 90 +++++ .../storage/db/repositories/test_user_logs.py | 64 +++ .../db/repositories/test_user_tools.py | 98 +++++ tests/storage/db/repositories/test_users.py | 222 +++++++++++ 36 files changed, 1929 insertions(+), 87 deletions(-) create mode 100644 application/storage/db/repositories/feedback.py create mode 100644 application/storage/db/repositories/prompts.py create mode 100644 application/storage/db/repositories/stack_logs.py create mode 100644 application/storage/db/repositories/token_usage.py create mode 100644 application/storage/db/repositories/user_logs.py create mode 100644 application/storage/db/repositories/user_tools.py create mode 100644 tests/storage/db/__init__.py create mode 100644 tests/storage/db/conftest.py create mode 100644 tests/storage/db/repositories/__init__.py create mode 100644 tests/storage/db/repositories/test_feedback.py create mode 100644 tests/storage/db/repositories/test_prompts.py create mode 100644 tests/storage/db/repositories/test_stack_logs.py create mode 100644 tests/storage/db/repositories/test_token_usage.py create mode 100644 tests/storage/db/repositories/test_user_logs.py create mode 100644 tests/storage/db/repositories/test_user_tools.py create mode 100644 tests/storage/db/repositories/test_users.py diff --git a/.github/styles/config/vocabularies/DocsGPT/accept.txt b/.github/styles/config/vocabularies/DocsGPT/accept.txt index 3d3de997..7605df49 100644 --- a/.github/styles/config/vocabularies/DocsGPT/accept.txt +++ b/.github/styles/config/vocabularies/DocsGPT/accept.txt @@ -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 diff --git a/application/agents/tools/brave.py b/application/agents/tools/brave.py index 66b21b10..8efc9197 100644 --- a/application/agents/tools/brave.py +++ b/application/agents/tools/brave.py @@ -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 { diff --git a/application/agents/tools/cryptoprice.py b/application/agents/tools/cryptoprice.py index c25c3d43..4dac5944 100644 --- a/application/agents/tools/cryptoprice.py +++ b/application/agents/tools/cryptoprice.py @@ -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: diff --git a/application/agents/tools/ntfy.py b/application/agents/tools/ntfy.py index 9a2d44ca..eac1b7ca 100644 --- a/application/agents/tools/ntfy.py +++ b/application/agents/tools/ntfy.py @@ -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): diff --git a/application/agents/tools/telegram.py b/application/agents/tools/telegram.py index d4381370..3007c300 100644 --- a/application/agents/tools/telegram.py +++ b/application/agents/tools/telegram.py @@ -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): diff --git a/application/api/answer/routes/base.py b/application/api/answer/routes/base.py index 4a152b2a..372b8cc2 100644 --- a/application/api/answer/routes/base.py +++ b/application/api/answer/routes/base.py @@ -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: diff --git a/application/api/user/prompts/routes.py b/application/api/user/prompts/routes.py index a08d70c8..e55c9d0e 100644 --- a/application/api/user/prompts/routes.py +++ b/application/api/user/prompts/routes.py @@ -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) diff --git a/application/api/user/tools/routes.py b/application/api/user/tools/routes.py index 23014845..0ccb371c 100644 --- a/application/api/user/tools/routes.py +++ b/application/api/user/tools/routes.py @@ -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 diff --git a/application/core/settings.py b/application/core/settings.py index c9d3cf84..2ed653cb 100644 --- a/application/core/settings.py +++ b/application/core/settings.py @@ -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 diff --git a/application/logging.py b/application/logging.py index 2c5cde27..1361b0f1 100644 --- a/application/logging.py +++ b/application/logging.py @@ -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) diff --git a/application/parser/connectors/share_point/loader.py b/application/parser/connectors/share_point/loader.py index 191e3e54..7e0ed6dd 100644 --- a/application/parser/connectors/share_point/loader.py +++ b/application/parser/connectors/share_point/loader.py @@ -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() diff --git a/application/parser/file/docs_parser.py b/application/parser/file/docs_parser.py index a1295290..2e5fa28d 100644 --- a/application/parser/file/docs_parser.py +++ b/application/parser/file/docs_parser.py @@ -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 diff --git a/application/parser/file/image_parser.py b/application/parser/file/image_parser.py index c7aa66a3..b6b3d775 100644 --- a/application/parser/file/image_parser.py +++ b/application/parser/file/image_parser.py @@ -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 = "" diff --git a/application/parser/remote/github_loader.py b/application/parser/remote/github_loader.py index 327e59a2..f51ddc18 100644 --- a/application/parser/remote/github_loader.py +++ b/application/parser/remote/github_loader.py @@ -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 diff --git a/application/storage/db/engine.py b/application/storage/db/engine.py index 723a1eb7..9ca0fff4 100644 --- a/application/storage/db/engine.py +++ b/application/storage/db/engine.py @@ -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 diff --git a/application/storage/db/models.py b/application/storage/db/models.py index 8b3d0527..dd87c616 100644 --- a/application/storage/db/models.py +++ b/application/storage/db/models.py @@ -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()), +) diff --git a/application/storage/db/repositories/feedback.py b/application/storage/db/repositories/feedback.py new file mode 100644 index 00000000..05d7704f --- /dev/null +++ b/application/storage/db/repositories/feedback.py @@ -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()] diff --git a/application/storage/db/repositories/prompts.py b/application/storage/db/repositories/prompts.py new file mode 100644 index 00000000..e959f33f --- /dev/null +++ b/application/storage/db/repositories/prompts.py @@ -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) diff --git a/application/storage/db/repositories/stack_logs.py b/application/storage/db/repositories/stack_logs.py new file mode 100644 index 00000000..eab3efce --- /dev/null +++ b/application/storage/db/repositories/stack_logs.py @@ -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, + }, + ) diff --git a/application/storage/db/repositories/token_usage.py b/application/storage/db/repositories/token_usage.py new file mode 100644 index 00000000..9171b65d --- /dev/null +++ b/application/storage/db/repositories/token_usage.py @@ -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() diff --git a/application/storage/db/repositories/user_logs.py b/application/storage/db/repositories/user_logs.py new file mode 100644 index 00000000..63f525d4 --- /dev/null +++ b/application/storage/db/repositories/user_logs.py @@ -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 diff --git a/application/storage/db/repositories/user_tools.py b/application/storage/db/repositories/user_tools.py new file mode 100644 index 00000000..7d8bfab8 --- /dev/null +++ b/application/storage/db/repositories/user_tools.py @@ -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 diff --git a/application/usage.py b/application/usage.py index e41a3588..f6b574c0 100644 --- a/application/usage.py +++ b/application/usage.py @@ -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): diff --git a/application/worker.py b/application/worker.py index 9821b65b..e98fc14d 100755 --- a/application/worker.py +++ b/application/worker.py @@ -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: diff --git a/docs/content/Guides/Architecture.mdx b/docs/content/Guides/Architecture.mdx index d981dc63..e902e1e7 100644 --- a/docs/content/Guides/Architecture.mdx +++ b/docs/content/Guides/Architecture.mdx @@ -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 diff --git a/scripts/db/backfill.py b/scripts/db/backfill.py index f08c8bf8..505bec3b 100644 --- a/scripts/db/backfill.py +++ b/scripts/db/backfill.py @@ -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, } diff --git a/tests/storage/db/__init__.py b/tests/storage/db/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tests/storage/db/conftest.py b/tests/storage/db/conftest.py new file mode 100644 index 00000000..8b765930 --- /dev/null +++ b/tests/storage/db/conftest.py @@ -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() diff --git a/tests/storage/db/repositories/__init__.py b/tests/storage/db/repositories/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tests/storage/db/repositories/test_feedback.py b/tests/storage/db/repositories/test_feedback.py new file mode 100644 index 00000000..78a99625 --- /dev/null +++ b/tests/storage/db/repositories/test_feedback.py @@ -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 diff --git a/tests/storage/db/repositories/test_prompts.py b/tests/storage/db/repositories/test_prompts.py new file mode 100644 index 00000000..aee9dbd9 --- /dev/null +++ b/tests/storage/db/repositories/test_prompts.py @@ -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"] diff --git a/tests/storage/db/repositories/test_stack_logs.py b/tests/storage/db/repositories/test_stack_logs.py new file mode 100644 index 00000000..0da1d8a0 --- /dev/null +++ b/tests/storage/db/repositories/test_stack_logs.py @@ -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 diff --git a/tests/storage/db/repositories/test_token_usage.py b/tests/storage/db/repositories/test_token_usage.py new file mode 100644 index 00000000..fc76a036 --- /dev/null +++ b/tests/storage/db/repositories/test_token_usage.py @@ -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 diff --git a/tests/storage/db/repositories/test_user_logs.py b/tests/storage/db/repositories/test_user_logs.py new file mode 100644 index 00000000..5f3af23b --- /dev/null +++ b/tests/storage/db/repositories/test_user_logs.py @@ -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" diff --git a/tests/storage/db/repositories/test_user_tools.py b/tests/storage/db/repositories/test_user_tools.py new file mode 100644 index 00000000..d3e36059 --- /dev/null +++ b/tests/storage/db/repositories/test_user_tools.py @@ -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 diff --git a/tests/storage/db/repositories/test_users.py b/tests/storage/db/repositories/test_users.py new file mode 100644 index 00000000..820d5ebf --- /dev/null +++ b/tests/storage/db/repositories/test_users.py @@ -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"] From 875868b7e5219472fe2cada1b5aef38ac1e90181 Mon Sep 17 00:00:00 2001 From: Alex Date: Sun, 12 Apr 2026 12:16:19 +0100 Subject: [PATCH 2/5] fix: comment --- application/core/settings.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/application/core/settings.py b/application/core/settings.py index 2ed653cb..4a5d1c49 100644 --- a/application/core/settings.py +++ b/application/core/settings.py @@ -28,9 +28,7 @@ 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. 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. + # User-data Postgres DB. POSTGRES_URI: Optional[str] = None # MongoDB→Postgres migration: dual-write to Postgres (Mongo stays source of truth) From 9388996a157fa997f411e0bff98995bc4a9bb35a Mon Sep 17 00:00:00 2001 From: Alex Date: Sun, 12 Apr 2026 12:18:31 +0100 Subject: [PATCH 3/5] fix: ruff --- tests/storage/db/conftest.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/storage/db/conftest.py b/tests/storage/db/conftest.py index 8b765930..0a8b9c47 100644 --- a/tests/storage/db/conftest.py +++ b/tests/storage/db/conftest.py @@ -19,7 +19,7 @@ import sys from pathlib import Path import pytest -from sqlalchemy import create_engine, text +from sqlalchemy import create_engine from application.core.settings import settings From a9761061fca310fd536a88be26f653c4869c8b5f Mon Sep 17 00:00:00 2001 From: Alex Date: Sun, 12 Apr 2026 12:24:58 +0100 Subject: [PATCH 4/5] fix: mini issues --- application/storage/db/engine.py | 37 ++++++------------ .../storage/db/repositories/user_tools.py | 39 +++++++++++++------ application/vectorstore/pgvector.py | 13 +++++-- .../db/repositories/test_user_tools.py | 6 ++- 4 files changed, 53 insertions(+), 42 deletions(-) diff --git a/application/storage/db/engine.py b/application/storage/db/engine.py index 9ca0fff4..5102b81e 100644 --- a/application/storage/db/engine.py +++ b/application/storage/db/engine.py @@ -24,38 +24,25 @@ _engine: Optional[Engine] = None def _resolve_uri() -> str: - """Pick the Postgres URI, falling back to the pgvector connection string. + """Return the Postgres URI for user-data tables. - 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. + 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. """ - 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'." - ) + 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'." + ) + return settings.POSTGRES_URI def get_engine() -> Engine: """Return the process-wide SQLAlchemy Engine, creating it if needed. - 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 Postgres via psycopg3. diff --git a/application/storage/db/repositories/user_tools.py b/application/storage/db/repositories/user_tools.py index 7d8bfab8..945347f6 100644 --- a/application/storage/db/repositories/user_tools.py +++ b/application/storage/db/repositories/user_tools.py @@ -78,18 +78,33 @@ class UserToolsRepository: 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) + params: dict = { + "id": tool_id, + "user_id": user_id, + "name": filtered.get("name"), + "custom_name": filtered.get("custom_name"), + "display_name": filtered.get("display_name"), + "config": ( + json.dumps(filtered["config"]) + if "config" in filtered and isinstance(filtered["config"], dict) + else filtered.get("config") + ), + } + self._conn.execute( + text( + """ + UPDATE user_tools + SET + name = COALESCE(:name, name), + custom_name = COALESCE(:custom_name, custom_name), + display_name = COALESCE(:display_name, display_name), + config = COALESCE(CAST(:config AS jsonb), config), + updated_at = now() + WHERE id = CAST(:id AS uuid) AND user_id = :user_id + """ + ), + params, + ) def delete(self, tool_id: str, user_id: str) -> bool: result = self._conn.execute( diff --git a/application/vectorstore/pgvector.py b/application/vectorstore/pgvector.py index eb4e7178..3066d8cd 100644 --- a/application/vectorstore/pgvector.py +++ b/application/vectorstore/pgvector.py @@ -27,13 +27,20 @@ class PGVectorStore(BaseVectorStore): self._metadata_column = metadata_column self._embedding = self._get_embeddings(settings.EMBEDDINGS_NAME, embeddings_key) - # Use provided connection string or fall back to settings + # Use provided connection string or fall back to settings. + # If PGVECTOR_CONNECTION_STRING is not set but POSTGRES_URI is, + # reuse the same cluster — normalize from SQLAlchemy dialect to libpq form. self._connection_string = connection_string or getattr(settings, 'PGVECTOR_CONNECTION_STRING', None) - + + if not self._connection_string and getattr(settings, 'POSTGRES_URI', None): + from application.core.db_uri import normalize_pgvector_connection_string + self._connection_string = normalize_pgvector_connection_string(settings.POSTGRES_URI) + if not self._connection_string: raise ValueError( "PostgreSQL connection string is required. " - "Set PGVECTOR_CONNECTION_STRING in settings or pass connection_string parameter." + "Set PGVECTOR_CONNECTION_STRING or POSTGRES_URI in settings, " + "or pass connection_string parameter." ) try: diff --git a/tests/storage/db/repositories/test_user_tools.py b/tests/storage/db/repositories/test_user_tools.py index d3e36059..9c830328 100644 --- a/tests/storage/db/repositories/test_user_tools.py +++ b/tests/storage/db/repositories/test_user_tools.py @@ -88,11 +88,13 @@ 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 + deleted = repo.delete(created["id"], "user-1") + assert deleted 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 + deleted = repo.delete(created["id"], "user-other") + assert deleted is False assert repo.get(created["id"]) is not None From b5b6538762096846928e5b0bca35f859b003b9e2 Mon Sep 17 00:00:00 2001 From: Alex Date: Sun, 12 Apr 2026 12:35:23 +0100 Subject: [PATCH 5/5] fix: tests --- tests/parser/remote/test_github_loader.py | 3 ++- tests/test_worker.py | 2 +- tests/vectorstore/test_pgvector.py | 1 + 3 files changed, 4 insertions(+), 2 deletions(-) diff --git a/tests/parser/remote/test_github_loader.py b/tests/parser/remote/test_github_loader.py index bd0a5fc0..32662279 100644 --- a/tests/parser/remote/test_github_loader.py +++ b/tests/parser/remote/test_github_loader.py @@ -31,6 +31,7 @@ class TestGitHubLoaderFetchFileContent: mock_get.assert_called_once_with( "https://api.github.com/repos/owner/repo/contents/README.md", headers=loader.headers, + timeout=100, ) @patch("application.parser.remote.github_loader.requests.get") @@ -66,7 +67,7 @@ class TestGitHubLoaderFetchRepoFiles: def test_recurses_directories(self, mock_get): loader = GitHubLoader() - def side_effect(url, headers=None): + def side_effect(url, headers=None, timeout=None): if url.endswith("/contents/"): return make_response([ {"type": "file", "path": "README.md"}, diff --git a/tests/test_worker.py b/tests/test_worker.py index 417ee160..29bb4f15 100644 --- a/tests/test_worker.py +++ b/tests/test_worker.py @@ -342,7 +342,7 @@ class TestDownloadFile: dest = str(tmp_path / "downloaded.txt") download_file("http://example.com/file", {"key": "val"}, dest) - mock_get.assert_called_once_with("http://example.com/file", params={"key": "val"}) + mock_get.assert_called_once_with("http://example.com/file", params={"key": "val"}, timeout=100) with open(dest, "rb") as f: assert f.read() == b"file content" diff --git a/tests/vectorstore/test_pgvector.py b/tests/vectorstore/test_pgvector.py index 08cfc11e..fa01279a 100644 --- a/tests/vectorstore/test_pgvector.py +++ b/tests/vectorstore/test_pgvector.py @@ -70,6 +70,7 @@ class TestPGVectorStoreInit: mock_get_emb.return_value = Mock(dimension=768) mock_settings.EMBEDDINGS_NAME = "test_model" mock_settings.PGVECTOR_CONNECTION_STRING = None + mock_settings.POSTGRES_URI = None from application.vectorstore.pgvector import PGVectorStore