crewplus 0.3.4__tar.gz → 0.3.5__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (80) hide show
  1. {crewplus-0.3.4 → crewplus-0.3.5}/PKG-INFO +1 -1
  2. {crewplus-0.3.4 → crewplus-0.3.5}/crewplus/vectorstores/milvus/vdb_service.py +162 -30
  3. {crewplus-0.3.4 → crewplus-0.3.5}/pyproject.toml +1 -1
  4. crewplus-0.3.5/tests/test_vdb_service_connection_reuse.py +189 -0
  5. {crewplus-0.3.4 → crewplus-0.3.5}/LICENSE +0 -0
  6. {crewplus-0.3.4 → crewplus-0.3.5}/README.md +0 -0
  7. {crewplus-0.3.4 → crewplus-0.3.5}/crewplus/__init__.py +0 -0
  8. {crewplus-0.3.4 → crewplus-0.3.5}/crewplus/__pycache__/__init__.cpython-311.pyc.1535298853824 +0 -0
  9. {crewplus-0.3.4 → crewplus-0.3.5}/crewplus/callbacks/__init__.py +0 -0
  10. {crewplus-0.3.4 → crewplus-0.3.5}/crewplus/callbacks/async_langfuse_handler.py +0 -0
  11. {crewplus-0.3.4 → crewplus-0.3.5}/crewplus/services/__init__.py +0 -0
  12. {crewplus-0.3.4 → crewplus-0.3.5}/crewplus/services/__pycache__/__init__.cpython-311.pyc.1535299230960 +0 -0
  13. {crewplus-0.3.4 → crewplus-0.3.5}/crewplus/services/__pycache__/gemini_chat_model.cpython-311.pyc.1535298773120 +0 -0
  14. {crewplus-0.3.4 → crewplus-0.3.5}/crewplus/services/__pycache__/runtime_model_factory.cpython-311.pyc.1959047557728 +0 -0
  15. {crewplus-0.3.4 → crewplus-0.3.5}/crewplus/services/azure_chat_model.py +0 -0
  16. {crewplus-0.3.4 → crewplus-0.3.5}/crewplus/services/capabilities/__init__.py +0 -0
  17. {crewplus-0.3.4 → crewplus-0.3.5}/crewplus/services/capabilities/base.py +0 -0
  18. {crewplus-0.3.4 → crewplus-0.3.5}/crewplus/services/capabilities/bash.py +0 -0
  19. {crewplus-0.3.4 → crewplus-0.3.5}/crewplus/services/capabilities/registry.py +0 -0
  20. {crewplus-0.3.4 → crewplus-0.3.5}/crewplus/services/capabilities/web_search.py +0 -0
  21. {crewplus-0.3.4 → crewplus-0.3.5}/crewplus/services/claude_chat_model.py +0 -0
  22. {crewplus-0.3.4 → crewplus-0.3.5}/crewplus/services/feedback.md +0 -0
  23. {crewplus-0.3.4 → crewplus-0.3.5}/crewplus/services/feedback_manager.py +0 -0
  24. {crewplus-0.3.4 → crewplus-0.3.5}/crewplus/services/gemini_chat_model.py +0 -0
  25. {crewplus-0.3.4 → crewplus-0.3.5}/crewplus/services/init_services.py +0 -0
  26. {crewplus-0.3.4 → crewplus-0.3.5}/crewplus/services/model_load_balancer.py +0 -0
  27. {crewplus-0.3.4 → crewplus-0.3.5}/crewplus/services/oai_claude_chat_model.py +0 -0
  28. {crewplus-0.3.4 → crewplus-0.3.5}/crewplus/services/runtime_model_factory.py +0 -0
  29. {crewplus-0.3.4 → crewplus-0.3.5}/crewplus/services/schemas/feedback.py +0 -0
  30. {crewplus-0.3.4 → crewplus-0.3.5}/crewplus/services/tracing_manager.py +0 -0
  31. {crewplus-0.3.4 → crewplus-0.3.5}/crewplus/utils/__init__.py +0 -0
  32. {crewplus-0.3.4 → crewplus-0.3.5}/crewplus/utils/schema_action.py +0 -0
  33. {crewplus-0.3.4 → crewplus-0.3.5}/crewplus/utils/schema_document_updater.py +0 -0
  34. {crewplus-0.3.4 → crewplus-0.3.5}/crewplus/utils/tracing_util.py +0 -0
  35. {crewplus-0.3.4 → crewplus-0.3.5}/crewplus/vectorstores/milvus/__init__.py +0 -0
  36. {crewplus-0.3.4 → crewplus-0.3.5}/crewplus/vectorstores/milvus/milvus_schema_manager.py +0 -0
  37. {crewplus-0.3.4 → crewplus-0.3.5}/crewplus/vectorstores/milvus/schema_milvus.py +0 -0
  38. {crewplus-0.3.4 → crewplus-0.3.5}/docs/GeminiChatModel.md +0 -0
  39. {crewplus-0.3.4 → crewplus-0.3.5}/docs/ModelLoadBalancer.md +0 -0
  40. {crewplus-0.3.4 → crewplus-0.3.5}/docs/VDBService.md +0 -0
  41. {crewplus-0.3.4 → crewplus-0.3.5}/docs/index.md +0 -0
  42. {crewplus-0.3.4 → crewplus-0.3.5}/tests/README.md +0 -0
  43. {crewplus-0.3.4 → crewplus-0.3.5}/tests/conftest.py +0 -0
  44. {crewplus-0.3.4 → crewplus-0.3.5}/tests/output/drawing_ingestion/gpt-5.4_run1.md +0 -0
  45. {crewplus-0.3.4 → crewplus-0.3.5}/tests/output/drawing_ingestion/gpt-5.4_run2.md +0 -0
  46. {crewplus-0.3.4 → crewplus-0.3.5}/tests/output/drawing_ingestion/gpt-5.6-luna_run1.md +0 -0
  47. {crewplus-0.3.4 → crewplus-0.3.5}/tests/output/drawing_ingestion/gpt-5.6-luna_run2.md +0 -0
  48. {crewplus-0.3.4 → crewplus-0.3.5}/tests/output/drawing_ingestion/gpt-5.6-terra_run1.md +0 -0
  49. {crewplus-0.3.4 → crewplus-0.3.5}/tests/output/drawing_ingestion/gpt-5.6-terra_run2.md +0 -0
  50. {crewplus-0.3.4 → crewplus-0.3.5}/tests/output/drawing_ingestion/page24_150dpi.png +0 -0
  51. {crewplus-0.3.4 → crewplus-0.3.5}/tests/output/drawing_ingestion/report.md +0 -0
  52. {crewplus-0.3.4 → crewplus-0.3.5}/tests/output/drawing_ingestion/results.json +0 -0
  53. {crewplus-0.3.4 → crewplus-0.3.5}/tests/output/drawing_ingestion/session1-luna-terra/gpt-5.6-luna_run1.md +0 -0
  54. {crewplus-0.3.4 → crewplus-0.3.5}/tests/output/drawing_ingestion/session1-luna-terra/gpt-5.6-luna_run2.md +0 -0
  55. {crewplus-0.3.4 → crewplus-0.3.5}/tests/output/drawing_ingestion/session1-luna-terra/gpt-5.6-terra_run1.md +0 -0
  56. {crewplus-0.3.4 → crewplus-0.3.5}/tests/output/drawing_ingestion/session1-luna-terra/gpt-5.6-terra_run2.md +0 -0
  57. {crewplus-0.3.4 → crewplus-0.3.5}/tests/output/drawing_ingestion/session1-luna-terra/report.md +0 -0
  58. {crewplus-0.3.4 → crewplus-0.3.5}/tests/output/drawing_ingestion/session1-luna-terra/results.json +0 -0
  59. {crewplus-0.3.4 → crewplus-0.3.5}/tests/output/drawing_ingestion/session2-55-sol/gpt-5.5_run1.md +0 -0
  60. {crewplus-0.3.4 → crewplus-0.3.5}/tests/output/drawing_ingestion/session2-55-sol/gpt-5.5_run2.md +0 -0
  61. {crewplus-0.3.4 → crewplus-0.3.5}/tests/output/drawing_ingestion/session2-55-sol/gpt-5.6-sol_run1.md +0 -0
  62. {crewplus-0.3.4 → crewplus-0.3.5}/tests/output/drawing_ingestion/session2-55-sol/gpt-5.6-sol_run2.md +0 -0
  63. {crewplus-0.3.4 → crewplus-0.3.5}/tests/output/drawing_ingestion/session2-55-sol/page24_150dpi.png +0 -0
  64. {crewplus-0.3.4 → crewplus-0.3.5}/tests/output/drawing_ingestion/session2-55-sol/report.md +0 -0
  65. {crewplus-0.3.4 → crewplus-0.3.5}/tests/output/drawing_ingestion/session2-55-sol/results.json +0 -0
  66. {crewplus-0.3.4 → crewplus-0.3.5}/tests/output/drawing_ingestion/session2-55-sol/summary-report-55-sol.md +0 -0
  67. {crewplus-0.3.4 → crewplus-0.3.5}/tests/output/drawing_ingestion/summary-report.md +0 -0
  68. {crewplus-0.3.4 → crewplus-0.3.5}/tests/test_bash_tools.py +0 -0
  69. {crewplus-0.3.4 → crewplus-0.3.5}/tests/test_claude_api.py +0 -0
  70. {crewplus-0.3.4 → crewplus-0.3.5}/tests/test_claude_custom_endpoint_integration.py +0 -0
  71. {crewplus-0.3.4 → crewplus-0.3.5}/tests/test_claude_langchain_openai.py +0 -0
  72. {crewplus-0.3.4 → crewplus-0.3.5}/tests/test_claude_taskflow_schema.py +0 -0
  73. {crewplus-0.3.4 → crewplus-0.3.5}/tests/test_custom_claude_endpoint.py +0 -0
  74. {crewplus-0.3.4 → crewplus-0.3.5}/tests/test_gemini3_pro_preview.py +0 -0
  75. {crewplus-0.3.4 → crewplus-0.3.5}/tests/test_gemini_bind_tools.py +0 -0
  76. {crewplus-0.3.4 → crewplus-0.3.5}/tests/test_oai_claude_chat_model.py +0 -0
  77. {crewplus-0.3.4 → crewplus-0.3.5}/tests/test_ops_models_streaming.py +0 -0
  78. {crewplus-0.3.4 → crewplus-0.3.5}/tests/test_pdf_drawing_ingestion_gpt55_sol.py +0 -0
  79. {crewplus-0.3.4 → crewplus-0.3.5}/tests/test_pdf_drawing_ingestion_gpt56.py +0 -0
  80. {crewplus-0.3.4 → crewplus-0.3.5}/tests/test_runtime_model_factory.py +0 -0
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.1
2
2
  Name: crewplus
3
- Version: 0.3.4
3
+ Version: 0.3.5
4
4
  Summary: Base services for CrewPlus AI applications
5
5
  Author-Email: Tim Liu <tim@opsmateai.com>
6
6
  License: MIT
@@ -5,7 +5,9 @@
5
5
  # @Last Modified time: 2025-10-09
6
6
 
7
7
  import asyncio
8
+ import hashlib
8
9
  import logging
10
+ import threading
9
11
  import time
10
12
  import uuid
11
13
  from collections import defaultdict
@@ -99,19 +101,42 @@ class VDBService(object):
99
101
  """
100
102
  _client: MilvusClient
101
103
  _async_client: Optional[AsyncMilvusClient] = None
104
+ # `_instances` / `_async_instances` 的 key 是 f"{conn_key}:{collection_name}",
105
+ # 不是裸 collection_name。历史上只用 collection_name 且本 dict 是类级共享的,
106
+ # 于是「同名 collection 分处两个 Zilliz 集群」时后来者会拿到前者的 store —— 数据被
107
+ # 静默写进别的租户的集群。加上 conn_key 前缀后该情形结构上不可能发生。
108
+ # 构造 key 一律走 self._instance_key(),不要手写拼接。
102
109
  _instances: Dict[str, Milvus] = {}
103
110
  _async_instances: Dict[str, Milvus] = {}
104
111
  _async_instance_locks: Dict[str, asyncio.Lock] = defaultdict(asyncio.Lock)
105
112
 
113
+ # 进程级 MilvusClient 缓存:同一组连接参数只建一条 gRPC 连接。
114
+ #
115
+ # 必须显式判缓存,不能只靠「固定别名」:pymilvus 的 connections.connect() 对已存在的
116
+ # 别名不做复用 —— 只校验地址是否一致,随后照样新建 GrpcHandler 并覆盖注册表,旧
117
+ # handler 既不 close 也不再可达。
118
+ #
119
+ # 为什么需要:VDBService.__init__ 每次都会真去建一条连接并注册进 pymilvus 的进程级
120
+ # 全局单例,而本包从不 disconnect;worker 用 --pool=threads,Celery 的
121
+ # worker_max_tasks_per_child 只支持 prefork,进程从不回收。于是「每任务一个
122
+ # VDBService」= 每任务一条新连接,批量摄取/删除时建连速率会打到每秒十几次,压垮网络
123
+ # 路径上的 per-pod/per-node 资源(2026-08-14 / 08-16 两次生产事故,单次建连失败率 57%)。
124
+ _clients: Dict[str, MilvusClient] = {}
125
+ # 每个 alias 一把创建锁:建连是网络 I/O(超时+退避最长 ~21s),不能让 A 集群的慢建连
126
+ # 阻塞 B 集群的首次建连。_clients_meta_lock 只保护锁字典本身(纯内存,微秒级)。
127
+ # 同一写法的先例:crewplus-search 的 _store_cache_key_locks。
128
+ _client_locks: Dict[str, threading.Lock] = {}
129
+ _clients_meta_lock = threading.Lock()
130
+
106
131
  schema: str
107
132
  embedding_function: Embeddings
108
133
  index_params: dict
109
134
  connection_args: dict
110
135
  settings: dict
111
136
 
112
- def __init__(self, settings: dict = None, endpoint: str = None, token: str = None, schema: str = None, logger: logging.Logger = None):
137
+ def __init__(self, settings: dict = None, endpoint: str = None, token: str = None, schema: str = None, logger: logging.Logger = None, reuse_connection: bool = True):
113
138
  """
114
- Initializes the VDBService.
139
+ Initializes the VDBService.
115
140
 
116
141
  Can be initialized in two ways:
117
142
  1. By providing a full `settings` dictionary for complex configurations.
@@ -126,6 +151,9 @@ class VDBService(object):
126
151
  token (str, optional): The token for authenticating with the Zilliz cluster.
127
152
  schema (str, optional): Default schema for new collections. Defaults to None.
128
153
  logger (logging.Logger, optional): Logger instance. Defaults to None.
154
+ reuse_connection (bool): 是否复用进程内已有的同参数连接。默认 True。
155
+ 传 False 恢复旧行为(本实例独占一条连接、不进进程缓存),仅作为出问题时
156
+ 无需回滚包版本的逃生口 —— 正常情况下不应有人传它。
129
157
  """
130
158
  self.logger = logger or logging.getLogger(__name__)
131
159
  self.collection_schema = None
@@ -162,9 +190,20 @@ class VDBService(object):
162
190
 
163
191
  self._provider = provider # Store provider for lazy initialization
164
192
 
165
- # Create separate aliases for sync and async clients to avoid connection handler race conditions.
166
- self.sync_alias = f"crewplus-vdb-sync-{uuid.uuid4()}"
167
- self.async_alias = f"crewplus-vdb-async-{uuid.uuid4()}"
193
+ # 别名从「每实例一个 uuid4」改为「按连接参数派生」。
194
+ #
195
+ # uuid4 别名让每个实例都往 pymilvus 的进程级注册表里塞一条独立连接,而且
196
+ # langchain_milvus 会把 connection_args['alias'] 原样传给它自己的 MilvusClient
197
+ # (langchain_milvus 0.2.1: self._milvus_client = MilvusClient(**connection_args)),
198
+ # 所以别名不稳定时 get_vector_store 也会各建各的连接。派生别名后,主 client 与所有
199
+ # langchain Milvus 实例落在同一个别名上,稳态下每个连接键只有一条活连接。
200
+ #
201
+ # 一个变量控制全部行为:reuse_connection=False 时 conn_key 取随机值 —— 别名与
202
+ # _instances 前缀的生成方式保持单一,逃生口不引入第二条代码路径。
203
+ self._reuse_connection = reuse_connection
204
+ self._conn_key = self._connection_key(provider) if reuse_connection else uuid.uuid4().hex
205
+ self.sync_alias = f"crewplus-vdb-sync-{self._conn_key}"
206
+ self.async_alias = f"crewplus-vdb-async-{self._conn_key}"
168
207
 
169
208
  # The default alias in connection_args should be the sync one, as langchain_milvus
170
209
  # primarily uses a synchronous client and will pick up this alias.
@@ -220,25 +259,79 @@ class VDBService(object):
220
259
  self.logger.error(f"Unsupported vector store provider: {provider}")
221
260
  raise NotImplementedError(f"Vector store provider '{provider}' is not supported.")
222
261
 
223
- def _initialize_milvus_client(self, provider: str) -> MilvusClient:
262
+ def _connection_key(self, provider: str) -> str:
263
+ """由连接身份派生的稳定 key:参数相同 -> 复用同一条连接。
264
+
265
+ token/password 必须纳入 —— 凭据不同必须走不同连接,否则跨租户串凭据。用 sha256
266
+ 截断而不是明文,是因为该 key 会被拼进 alias,而 alias 会出现在 INFO 日志里。
267
+ channel_options 刻意不纳入:它是常量(gRPC keepalive 参数),只会让 key 变长。
224
268
  """
225
- Initializes and returns a MilvusClient with a retry mechanism.
269
+ args = self._get_milvus_client_args(provider)
270
+ parts = [
271
+ provider,
272
+ str(args.get("uri") or ""),
273
+ str(args.get("host") or ""),
274
+ str(args.get("port") or ""),
275
+ str(args.get("db_name") or ""),
276
+ str(args.get("user") or ""),
277
+ str(args.get("token") or args.get("password") or ""),
278
+ ]
279
+ return hashlib.sha256("|".join(parts).encode("utf-8")).hexdigest()[:32]
280
+
281
+ def _instance_key(self, collection_name: str) -> str:
282
+ """`_instances` / `_async_instances` 的缓存 key。
283
+
284
+ 带 conn_key 前缀,让同名 collection 在不同集群/不同凭据下互不串用。
226
285
  """
286
+ return f"{self._conn_key}:{collection_name}"
287
+
288
+ def _create_milvus_client(self, provider: str, alias: str) -> MilvusClient:
289
+ """真正建连(含一次带退避的重试)。是否缓存由调用方决定。"""
227
290
  client_args = self._get_milvus_client_args(provider)
228
- client_args["alias"] = self.sync_alias
291
+ client_args["alias"] = alias
229
292
 
230
293
  try:
231
294
  # First attempt to connect
232
295
  return MilvusClient(**client_args)
233
296
  except Exception as e:
234
297
  self.logger.error(f"Failed to initialize MilvusClient, trying again. Error: {e}")
235
- # Second attempt after failure
298
+ # 退避后再试。原来是紧贴着重打,实测第二次成功率仅 9.6%(08-16:2105 次首次
299
+ # 失败里只有 201 次被第二次救回),却把建连总量放大了 30% —— 在一个「建连速率
300
+ # 就是病因」的故障里等于浇油。与异步版保持一致,睡 1 秒。
301
+ time.sleep(1)
236
302
  try:
237
303
  return MilvusClient(**client_args)
238
304
  except Exception as e_retry:
239
305
  self.logger.error(f"Failed to initialize MilvusClient on retry. Final error: {e_retry}")
240
306
  raise RuntimeError(f"Could not initialize MilvusClient after retry: {e_retry}")
241
307
 
308
+ @classmethod
309
+ def _get_client_lock(cls, alias: str) -> threading.Lock:
310
+ with cls._clients_meta_lock:
311
+ return cls._client_locks.setdefault(alias, threading.Lock())
312
+
313
+ def _initialize_milvus_client(self, provider: str) -> MilvusClient:
314
+ """返回本进程内该连接参数对应的 MilvusClient,没有则建一条并缓存。"""
315
+ alias = self.sync_alias
316
+
317
+ if not self._reuse_connection:
318
+ # 逃生口:旧行为。刻意不写入 _clients —— 每实例一条缓存条目才是真泄漏。
319
+ return self._create_milvus_client(provider, alias)
320
+
321
+ cached = VDBService._clients.get(alias) # 快路径:CPython 下 dict.get 是原子操作
322
+ if cached is not None:
323
+ return cached
324
+
325
+ with VDBService._get_client_lock(alias): # 只串行化同一 alias 的并发首建
326
+ cached = VDBService._clients.get(alias) # 双检:并发涌入的线程只让一个真去建
327
+ if cached is not None:
328
+ return cached
329
+
330
+ client = self._create_milvus_client(provider, alias)
331
+ VDBService._clients[alias] = client
332
+ self.logger.info(f"MilvusClient created and cached, alias={alias}")
333
+ return client
334
+
242
335
  def _initialize_async_milvus_client(self, provider: str) -> AsyncMilvusClient:
243
336
  """
244
337
  Initializes and returns an AsyncMilvusClient with a retry mechanism.
@@ -283,6 +376,34 @@ class VDBService(object):
283
376
  self._async_client = AsyncMilvusClient(**client_args)
284
377
  return self._async_client
285
378
 
379
+ def close(self) -> None:
380
+ """作废本连接:移出进程缓存、断开、并清掉挂在它上面的 langchain 实例。
381
+
382
+ 只在确认连接不可用时调用。这条连接是进程内共享的,关掉会影响所有正在用同一
383
+ endpoint+凭据的调用方 —— 它们会在下次构造 VDBService 时自动重建,所以是安全的,
384
+ 代价只是多一次握手。
385
+
386
+ 为什么需要它:`_instances` 历史上只在 drop_collection 里被删过,没有健康检查也没有
387
+ TTL,一条坏掉的 langchain Milvus 实例会被后续所有任务无条件复用,把整个 pod 的剩余
388
+ 生命周期拖死。这是既有缺陷(与连接复用无关),此处一并补上补救路径。
389
+ """
390
+ with VDBService._get_client_lock(self.sync_alias):
391
+ VDBService._clients.pop(self.sync_alias, None)
392
+ try:
393
+ connections.disconnect(self.sync_alias)
394
+ except Exception as e:
395
+ self.logger.warning(f"disconnect({self.sync_alias}) failed: {e}")
396
+
397
+ # 只清本连接键下的实例,别动其他租户的
398
+ prefix = f"{self._conn_key}:"
399
+ for cache in (VDBService._instances, VDBService._async_instances):
400
+ for k in [k for k in cache if k.startswith(prefix)]:
401
+ cache.pop(k, None)
402
+
403
+ self._client = None
404
+ self._async_client = None
405
+ self.logger.info(f"VDBService connection closed and evicted, alias={self.sync_alias}")
406
+
286
407
  def get_vector_field(self, collection_name: str) -> str:
287
408
  """
288
409
  Retrieves the vector field name for a given collection from a cached instance.
@@ -296,10 +417,11 @@ class VDBService(object):
296
417
  Raises:
297
418
  ValueError: If no cached instance is found for the collection.
298
419
  """
299
- if collection_name in self._instances:
300
- return self._instances[collection_name]._vector_field
301
- if collection_name in self._async_instances:
302
- return self._async_instances[collection_name]._vector_field
420
+ cache_key = self._instance_key(collection_name)
421
+ if cache_key in self._instances:
422
+ return self._instances[cache_key]._vector_field
423
+ if cache_key in self._async_instances:
424
+ return self._async_instances[cache_key]._vector_field
303
425
 
304
426
  self.logger.warning(f"No cached instance found for collection '{collection_name}' to get vector field. Creating a temporary sync instance.")
305
427
  # As a fallback, create a temporary sync instance to fetch the schema info.
@@ -516,9 +638,10 @@ class VDBService(object):
516
638
  raise ValueError("collection_name must be provided.")
517
639
 
518
640
  # Check for a cached instance. If found, return it immediately.
519
- if collection_name in self._instances:
641
+ cache_key = self._instance_key(collection_name)
642
+ if cache_key in self._instances:
520
643
  self.logger.info(f"Returning existing vector store instance for collection: {collection_name}")
521
- return self._instances[collection_name]
644
+ return self._instances[cache_key]
522
645
 
523
646
  self.logger.info(f"Creating new vector store instance for collection: {collection_name}")
524
647
  if embeddings is None:
@@ -555,7 +678,7 @@ class VDBService(object):
555
678
  )
556
679
 
557
680
  # Cache the newly created instance.
558
- self._instances[collection_name] = vdb
681
+ self._instances[cache_key] = vdb
559
682
 
560
683
  return vdb
561
684
 
@@ -584,11 +707,13 @@ class VDBService(object):
584
707
  self.logger.error("aget_vector_store called with no collection_name.")
585
708
  raise ValueError("collection_name must be provided.")
586
709
 
587
- lock = self._async_instance_locks[collection_name]
710
+ # 锁也按 conn_key 分片:不同集群的同名 collection 不该互相阻塞。
711
+ cache_key = self._instance_key(collection_name)
712
+ lock = self._async_instance_locks[cache_key]
588
713
  async with lock:
589
- if collection_name in self._async_instances:
714
+ if cache_key in self._async_instances:
590
715
  self.logger.info(f"Returning existing async vector store instance for collection: {collection_name} (post-lock)")
591
- return self._async_instances[collection_name]
716
+ return self._async_instances[cache_key]
592
717
 
593
718
  self.logger.info(f"Creating new async vector store instance for collection: {collection_name}")
594
719
  if embeddings is None:
@@ -612,7 +737,7 @@ class VDBService(object):
612
737
  self.logger.info(f"Swapping to async alias for instance of collection {collection_name}")
613
738
  vdb.aclient._using = self.async_alias
614
739
 
615
- self._async_instances[collection_name] = vdb
740
+ self._async_instances[cache_key] = vdb
616
741
  return vdb
617
742
 
618
743
  async def _acreate_milvus_instance_with_retry(
@@ -638,12 +763,17 @@ class VDBService(object):
638
763
  # --- START: CRITICAL FIX ---
639
764
  # Manually connect within the thread before creating the Milvus instance.
640
765
  # This ensures pymilvus registers the connection details for the current thread.
641
- try:
642
- connections.connect(**conn_args)
643
- self.logger.info(f"Successfully connected to Milvus with alias '{self.sync_alias}' in thread.")
644
- except Exception as e:
645
- self.logger.error(f"Failed to manually connect in thread: {e}")
646
- raise
766
+ #
767
+ # 已有同名连接就直接复用:pymilvus 的 connect() 对已存在的别名不做复用,
768
+ # 只校验地址一致,随后照样新建 GrpcHandler 并覆盖注册表 —— 会把
769
+ # _clients 里缓存的那条主连接的 handler 顶掉(旧的既不 close 也不再可达)。
770
+ if not connections.has_connection(self.sync_alias):
771
+ try:
772
+ connections.connect(**conn_args)
773
+ self.logger.info(f"Successfully connected to Milvus with alias '{self.sync_alias}' in thread.")
774
+ except Exception as e:
775
+ self.logger.error(f"Failed to manually connect in thread: {e}")
776
+ raise
647
777
 
648
778
  # Now, creating the Milvus instance will find the existing connection via the alias.
649
779
  instance = Milvus(
@@ -729,8 +859,9 @@ class VDBService(object):
729
859
  raise RuntimeError(f"An error occurred while dropping collection '{collection_name}'.") from e
730
860
  finally:
731
861
  # Whether successful or not, remove the stale instance from the cache.
732
- if collection_name in self._instances:
733
- del self._instances[collection_name]
862
+ cache_key = self._instance_key(collection_name)
863
+ if cache_key in self._instances:
864
+ del self._instances[cache_key]
734
865
  self.logger.info(f"Removed '{collection_name}' from instance cache.")
735
866
 
736
867
  async def adrop_collection(self, collection_name: str) -> None:
@@ -759,8 +890,9 @@ class VDBService(object):
759
890
  raise RuntimeError(f"An error occurred while dropping collection '{collection_name}' asynchronously.") from e
760
891
  finally:
761
892
  # Whether successful or not, remove the stale instance from the cache.
762
- if collection_name in self._async_instances:
763
- del self._async_instances[collection_name]
893
+ cache_key = self._instance_key(collection_name)
894
+ if cache_key in self._async_instances:
895
+ del self._async_instances[cache_key]
764
896
  self.logger.info(f"Removed '{collection_name}' from instance cache.")
765
897
 
766
898
  def delete_data_by_filter(self, collection_name: str = None, filter: str = None) -> None:
@@ -6,7 +6,7 @@ build-backend = "pdm.backend"
6
6
 
7
7
  [project]
8
8
  name = "crewplus"
9
- version = "0.3.4"
9
+ version = "0.3.5"
10
10
  description = "Base services for CrewPlus AI applications"
11
11
  authors = [
12
12
  { name = "Tim Liu", email = "tim@opsmateai.com" },
@@ -0,0 +1,189 @@
1
+ # -*- coding: utf-8 -*-
2
+ """VDBService 连接复用与缓存隔离的回归测试。
3
+
4
+ 背景:2026-08-14 / 08-16 两次生产事故。VDBService.__init__ 每次都会真去建一条 gRPC 连接并
5
+ 注册进 pymilvus 的进程级全局单例,而本包从不 disconnect;worker 用 --pool=threads,Celery 的
6
+ worker_max_tasks_per_child 只支持 prefork,进程从不回收。结果「每任务一个 VDBService」= 每
7
+ 任务一条新连接,批量摄取/删除时建连速率打到每秒十几次,单次建连失败率 57%。
8
+
9
+ 这些测试锁住三件事,防止改回去:
10
+ 1. 同参数构造 N 次只建 1 条连接;
11
+ 2. 不同 endpoint / 不同 token 绝不共享连接(凭据隔离);
12
+ 3. `_instances` 的 key 带连接身份 —— 同名 collection 分处两个集群时不会串数据。
13
+
14
+ 全部用 mock,不碰真实 Milvus。
15
+ """
16
+ import threading
17
+ from unittest.mock import MagicMock, patch
18
+
19
+ import pytest
20
+
21
+ from crewplus.vectorstores.milvus.vdb_service import VDBService
22
+
23
+
24
+ ENDPOINT_A = "https://in01-aaaaaaaaaaaaaaa.az-eastus2.vectordb.zillizcloud.com:19530"
25
+ ENDPOINT_B = "https://in01-bbbbbbbbbbbbbbb.az-eastus2.vectordb.zillizcloud.com:19530"
26
+ TOKEN_A = "token-tenant-a"
27
+ TOKEN_B = "token-tenant-b"
28
+
29
+
30
+ @pytest.fixture(autouse=True)
31
+ def _clean_process_caches():
32
+ """每个用例前后清空类级缓存 —— 它们是进程级的,会跨用例串。"""
33
+ VDBService._clients.clear()
34
+ VDBService._instances.clear()
35
+ VDBService._async_instances.clear()
36
+ yield
37
+ VDBService._clients.clear()
38
+ VDBService._instances.clear()
39
+ VDBService._async_instances.clear()
40
+
41
+
42
+ @pytest.fixture
43
+ def milvus_client_factory():
44
+ """替换 MilvusClient,返回每次调用都不同的 mock,并记录调用次数。"""
45
+ with patch("crewplus.vectorstores.milvus.vdb_service.MilvusClient") as m:
46
+ m.side_effect = lambda **kwargs: MagicMock(name=f"client-{m.call_count}")
47
+ yield m
48
+
49
+
50
+ def test_same_params_build_one_connection(milvus_client_factory):
51
+ """同一组 endpoint+token 构造 50 次,只应真建 1 条连接。"""
52
+ services = [VDBService(endpoint=ENDPOINT_A, token=TOKEN_A) for _ in range(50)]
53
+
54
+ assert milvus_client_factory.call_count == 1, (
55
+ "每次构造都建连 = 事故复现。连接数必须与构造次数解耦。"
56
+ )
57
+ # 50 个实例拿到的是同一个 client 对象
58
+ assert len({id(s._client) for s in services}) == 1
59
+ # 别名也必须稳定,否则 langchain_milvus 会照着新别名再建一条
60
+ assert len({s.sync_alias for s in services}) == 1
61
+ assert len(VDBService._clients) == 1
62
+
63
+
64
+ def test_different_endpoint_never_shares_connection(milvus_client_factory):
65
+ """不同集群必须各自建连。"""
66
+ a = VDBService(endpoint=ENDPOINT_A, token=TOKEN_A)
67
+ b = VDBService(endpoint=ENDPOINT_B, token=TOKEN_A)
68
+
69
+ assert milvus_client_factory.call_count == 2
70
+ assert a.sync_alias != b.sync_alias
71
+ assert a._client is not b._client
72
+
73
+
74
+ def test_different_token_never_shares_connection(milvus_client_factory):
75
+ """同集群但凭据不同也必须各自建连 —— 否则跨租户串凭据。"""
76
+ a = VDBService(endpoint=ENDPOINT_A, token=TOKEN_A)
77
+ b = VDBService(endpoint=ENDPOINT_A, token=TOKEN_B)
78
+
79
+ assert milvus_client_factory.call_count == 2, (
80
+ "token 必须参与连接键,否则 B 的写入会以 A 的凭据执行。"
81
+ )
82
+ assert a.sync_alias != b.sync_alias
83
+
84
+
85
+ def test_token_not_leaked_into_alias(milvus_client_factory):
86
+ """别名会进 INFO 日志,明文 token 不许出现在里面。"""
87
+ svc = VDBService(endpoint=ENDPOINT_A, token=TOKEN_A)
88
+ assert TOKEN_A not in svc.sync_alias
89
+ assert TOKEN_A not in svc.async_alias
90
+
91
+
92
+ def test_instance_cache_key_isolates_same_collection_across_clusters(milvus_client_factory):
93
+ """同名 collection 分处两个集群时,_instances 不许串。
94
+
95
+ 这是数据安全项:串了会把一个租户的文件静默写进另一个租户的向量库。
96
+ """
97
+ a = VDBService(endpoint=ENDPOINT_A, token=TOKEN_A)
98
+ b = VDBService(endpoint=ENDPOINT_B, token=TOKEN_B)
99
+
100
+ collection = "cp_shared_name_2026081700ab" # 迁移/恢复场景下两边会同名
101
+ store_a, store_b = MagicMock(name="store-a"), MagicMock(name="store-b")
102
+
103
+ with patch.object(VDBService, "_check_collection_exists", return_value=True), \
104
+ patch.object(VDBService, "_create_milvus_instance_with_retry", side_effect=[store_a, store_b]):
105
+ got_a = a.get_vector_store(collection, embeddings=MagicMock())
106
+ got_b = b.get_vector_store(collection, embeddings=MagicMock())
107
+
108
+ assert got_a is store_a
109
+ assert got_b is store_b, (
110
+ "B 拿到了 A 的 store —— 数据会被写进 A 的集群。缓存 key 必须带连接身份。"
111
+ )
112
+ assert len(VDBService._instances) == 2
113
+
114
+
115
+ def test_same_cluster_same_collection_shares_instance(milvus_client_factory):
116
+ """同集群同 collection 仍应复用 —— 这才是缓存的本意。"""
117
+ a = VDBService(endpoint=ENDPOINT_A, token=TOKEN_A)
118
+ b = VDBService(endpoint=ENDPOINT_A, token=TOKEN_A)
119
+
120
+ collection = "cp_kb_2026081700ab"
121
+ store = MagicMock(name="store")
122
+
123
+ with patch.object(VDBService, "_check_collection_exists", return_value=True), \
124
+ patch.object(VDBService, "_create_milvus_instance_with_retry", return_value=store):
125
+ assert a.get_vector_store(collection, embeddings=MagicMock()) is store
126
+ assert b.get_vector_store(collection, embeddings=MagicMock()) is store
127
+
128
+ assert len(VDBService._instances) == 1
129
+
130
+
131
+ def test_reuse_connection_false_restores_old_behavior(milvus_client_factory):
132
+ """逃生口:传 False 恢复独占连接,且不污染进程缓存。"""
133
+ a = VDBService(endpoint=ENDPOINT_A, token=TOKEN_A, reuse_connection=False)
134
+ b = VDBService(endpoint=ENDPOINT_A, token=TOKEN_A, reuse_connection=False)
135
+
136
+ assert milvus_client_factory.call_count == 2
137
+ assert a.sync_alias != b.sync_alias
138
+ assert VDBService._clients == {}, "退出口不得写入进程缓存,否则每实例一条才是真泄漏。"
139
+
140
+
141
+ def test_close_evicts_client_and_instances(milvus_client_factory):
142
+ """close() 必须同时清掉客户端缓存和挂在该连接上的 langchain 实例。"""
143
+ svc = VDBService(endpoint=ENDPOINT_A, token=TOKEN_A)
144
+ other = VDBService(endpoint=ENDPOINT_B, token=TOKEN_B)
145
+
146
+ VDBService._instances[svc._instance_key("cp_x")] = MagicMock()
147
+ VDBService._instances[other._instance_key("cp_x")] = MagicMock()
148
+
149
+ with patch("crewplus.vectorstores.milvus.vdb_service.connections") as conns:
150
+ svc.close()
151
+ conns.disconnect.assert_called_once_with(svc.sync_alias)
152
+
153
+ assert svc.sync_alias not in VDBService._clients
154
+ assert svc._instance_key("cp_x") not in VDBService._instances
155
+ # 别动别的租户
156
+ assert other._instance_key("cp_x") in VDBService._instances
157
+ assert other.sync_alias in VDBService._clients
158
+
159
+
160
+ def test_close_then_reconstruct_rebuilds(milvus_client_factory):
161
+ """close() 之后再构造应重新建连(失效路径可恢复)。"""
162
+ VDBService(endpoint=ENDPOINT_A, token=TOKEN_A).close()
163
+ VDBService(endpoint=ENDPOINT_A, token=TOKEN_A)
164
+ assert milvus_client_factory.call_count == 2
165
+
166
+
167
+ def test_concurrent_construction_builds_one_connection(milvus_client_factory):
168
+ """50 线程并发构造(worker 是 --pool=threads --concurrency=50)只应建 1 条连接。"""
169
+ barrier = threading.Barrier(50)
170
+ results, errors = [], []
171
+
172
+ def worker():
173
+ try:
174
+ barrier.wait(timeout=10)
175
+ results.append(VDBService(endpoint=ENDPOINT_A, token=TOKEN_A))
176
+ except Exception as e: # pragma: no cover
177
+ errors.append(e)
178
+
179
+ threads = [threading.Thread(target=worker) for _ in range(50)]
180
+ for t in threads:
181
+ t.start()
182
+ for t in threads:
183
+ t.join(timeout=20)
184
+
185
+ assert not errors, f"并发构造抛错: {errors}"
186
+ assert len(results) == 50
187
+ assert milvus_client_factory.call_count == 1, (
188
+ "双检锁失效:并发下建了多条连接。"
189
+ )
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes