Skip to content

Commit 69354ba

Browse files
authored
Merge branch 'arc53:main' into multimodal-clean
2 parents 6927fbf + 08822c3 commit 69354ba

12 files changed

Lines changed: 1980 additions & 1547 deletions

File tree

application/requirements.txt

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -13,12 +13,12 @@ onnxruntime>=1.19.0
1313
docx2txt==0.9
1414
ddgs>=8.0.0
1515
fast-ebook
16-
elevenlabs==2.41.0
16+
elevenlabs==2.43.0
1717
Flask==3.1.3
1818
faiss-cpu==1.13.2
19-
fastmcp==3.2.0
19+
fastmcp==3.2.4
2020
flask-restx==1.3.2
21-
google-genai==1.69.0
21+
google-genai==1.73.1
2222
google-api-python-client==2.193.0
2323
google-auth-httplib2==0.3.1
2424
google-auth-oauthlib==1.3.1
@@ -47,7 +47,7 @@ msal==1.35.1
4747
mypy-extensions==1.1.0
4848
networkx==3.6.1
4949
numpy==2.4.4
50-
openai==2.30.0
50+
openai==2.32.0
5151
openapi3-parser==1.1.22
5252
orjson==3.11.7
5353
packaging==26.0

application/storage/db/engine.py

Lines changed: 18 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -16,7 +16,7 @@
1616

1717
from typing import Optional
1818

19-
from sqlalchemy import Engine, create_engine
19+
from sqlalchemy import Engine, create_engine, event
2020

2121
from application.core.settings import settings
2222

@@ -43,17 +43,17 @@ def _resolve_uri() -> str:
4343
#: Per-statement wall-clock cap applied to every connection handed out by
4444
#: the engine. 30s is generous for interactive hot paths (reads under a few
4545
#: hundred ms are normal) but still catches a runaway query before it
46-
#: stacks up on PgBouncer or holds locks indefinitely. Override by
47-
#: rebuilding the engine with a different ``connect_args`` in tests.
46+
#: stacks up on PgBouncer or holds locks indefinitely.
4847
STATEMENT_TIMEOUT_MS = 30_000
4948

5049

5150
def get_engine() -> Engine:
5251
"""Return the process-wide SQLAlchemy Engine, creating it if needed.
5352
5453
The engine applies a server-side ``statement_timeout`` to every
55-
connection it hands out, so both :func:`db_session` and
56-
:func:`db_readonly` inherit the same guardrail.
54+
connection it hands out via a ``connect`` event, so both
55+
:func:`db_session` and :func:`db_readonly` inherit the same
56+
guardrail.
5757
5858
Returns:
5959
A SQLAlchemy ``Engine`` configured with a pooled connection to
@@ -68,13 +68,20 @@ def get_engine() -> Engine:
6868
pool_pre_ping=True, # survive PgBouncer / idle-disconnect recycles
6969
pool_recycle=1800,
7070
future=True,
71-
connect_args={
72-
# ``-c`` passes a GUC to the backend at connect time. This
73-
# covers *all* sessions — interactive, Celery, seeder — so
74-
# no route-handler can opt out by accident.
75-
"options": f"-c statement_timeout={STATEMENT_TIMEOUT_MS}",
76-
},
7771
)
72+
73+
@event.listens_for(_engine, "connect")
74+
def _apply_session_guardrails(dbapi_conn, _record):
75+
# Apply as a SQL ``SET`` (not a libpq ``options=-c ...``
76+
# startup parameter) so the engine works behind
77+
# PgBouncer-style poolers — notably Neon's ``-pooler``
78+
# endpoint, which rejects startup options. Explicit
79+
# ``commit()`` so the session-level SET survives SA's
80+
# transaction resets on pool return.
81+
with dbapi_conn.cursor() as cur:
82+
cur.execute(f"SET statement_timeout = {STATEMENT_TIMEOUT_MS}")
83+
dbapi_conn.commit()
84+
7885
return _engine
7986

8087

application/vectorstore/mongodb.py

Lines changed: 28 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -1,9 +1,23 @@
11
import logging
2+
from functools import cached_property
3+
24
from application.core.settings import settings
35
from application.vectorstore.base import BaseVectorStore
46
from application.vectorstore.document_class import Document
57

68

9+
def _lazy_import_pymongo():
10+
"""Lazy import of pymongo so installations that don't use the MongoDB vectorstore don't need it."""
11+
try:
12+
import pymongo
13+
except ImportError as exc:
14+
raise ImportError(
15+
"Could not import pymongo python package. "
16+
"Please install it with `pip install pymongo`."
17+
) from exc
18+
return pymongo
19+
20+
721
class MongoDBVectorStore(BaseVectorStore):
822
def __init__(
923
self,
@@ -20,20 +34,23 @@ def __init__(
2034
self._embedding_key = embedding_key
2135
self._embeddings_key = embeddings_key
2236
self._mongo_uri = settings.MONGO_URI
37+
self._database_name = database
38+
self._collection_name = collection
2339
self._source_id = source_id.replace("application/indexes/", "").rstrip("/")
2440
self._embedding = self._get_embeddings(settings.EMBEDDINGS_NAME, embeddings_key)
2541

26-
try:
27-
import pymongo
28-
except ImportError:
29-
raise ImportError(
30-
"Could not import pymongo python package. "
31-
"Please install it with `pip install pymongo`."
32-
)
33-
34-
self._client = pymongo.MongoClient(self._mongo_uri)
35-
self._database = self._client[database]
36-
self._collection = self._database[collection]
42+
@cached_property
43+
def _client(self):
44+
pymongo = _lazy_import_pymongo()
45+
return pymongo.MongoClient(self._mongo_uri)
46+
47+
@cached_property
48+
def _database(self):
49+
return self._client[self._database_name]
50+
51+
@cached_property
52+
def _collection(self):
53+
return self._database[self._collection_name]
3754

3855
def search(self, question, k=2, *args, **kwargs):
3956
query_vector = self._embedding.embed_query(question)

0 commit comments

Comments
 (0)