Compare commits
1
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
796f30e504 |
No files matched your search
@@ -16,11 +16,10 @@ Agent Memory V2 - VectorStore (Qdrant)
|
|||||||
- 简化运维,后续可按需拆分
|
- 简化运维,后续可按需拆分
|
||||||
"""
|
"""
|
||||||
|
|
||||||
import asyncio
|
|
||||||
import concurrent.futures
|
import concurrent.futures
|
||||||
import os
|
import os
|
||||||
import threading
|
import threading
|
||||||
from typing import Optional, List, Dict, Any, Tuple
|
from typing import Optional, List, Dict, Any, Tuple, Callable
|
||||||
from pathlib import Path
|
from pathlib import Path
|
||||||
import logging
|
import logging
|
||||||
|
|
||||||
@@ -30,12 +29,10 @@ from ..models.memory import (
|
|||||||
from ..config import MemoryConfig
|
from ..config import MemoryConfig
|
||||||
from ..utils.async_executor import run_in_executor_guarded
|
from ..utils.async_executor import run_in_executor_guarded
|
||||||
from .vector_store_base import VectorStoreBase
|
from .vector_store_base import VectorStoreBase
|
||||||
|
from .vector_store_chroma import _vdb_executor
|
||||||
|
|
||||||
logger = logging.getLogger(__name__)
|
logger = logging.getLogger(__name__)
|
||||||
|
|
||||||
# VDB 独立线程池(与 Chroma 共享同一策略)
|
|
||||||
from .vector_store_chroma import _vdb_executor
|
|
||||||
|
|
||||||
|
|
||||||
def _run_in_vdb_pool(func, *args, **kwargs):
|
def _run_in_vdb_pool(func, *args, **kwargs):
|
||||||
"""在 VDB 独立线程池中执行同步函数"""
|
"""在 VDB 独立线程池中执行同步函数"""
|
||||||
@@ -44,6 +41,70 @@ def _run_in_vdb_pool(func, *args, **kwargs):
|
|||||||
return run_in_executor_guarded(_vdb_executor, call)
|
return run_in_executor_guarded(_vdb_executor, call)
|
||||||
|
|
||||||
|
|
||||||
|
class _LocalQdrantOwner:
|
||||||
|
"""Keep one local Qdrant client and its SQLite state on one owner thread.
|
||||||
|
|
||||||
|
Local Qdrant may enable SQLite's ``check_same_thread`` on environments
|
||||||
|
compiled with ``THREADSAFE=2``. The public async surface still uses the
|
||||||
|
shared VDB pool, but every actual client access is forwarded to this
|
||||||
|
instance's single owner thread.
|
||||||
|
"""
|
||||||
|
|
||||||
|
def __init__(self, factory: Callable[[], Any]):
|
||||||
|
self._executor = concurrent.futures.ThreadPoolExecutor(
|
||||||
|
max_workers=1,
|
||||||
|
thread_name_prefix="qdrant-local-owner",
|
||||||
|
)
|
||||||
|
self._lifecycle_lock = threading.Lock()
|
||||||
|
self._closed = False
|
||||||
|
try:
|
||||||
|
self._client = self._executor.submit(factory).result()
|
||||||
|
except BaseException:
|
||||||
|
self._closed = True
|
||||||
|
self._executor.shutdown(wait=True, cancel_futures=True)
|
||||||
|
raise
|
||||||
|
|
||||||
|
def _submit(self, callback: Callable[[], Any]) -> Any:
|
||||||
|
with self._lifecycle_lock:
|
||||||
|
if self._closed:
|
||||||
|
raise RuntimeError("Local Qdrant client is closed")
|
||||||
|
future = self._executor.submit(callback)
|
||||||
|
return future.result()
|
||||||
|
|
||||||
|
def __getattr__(self, name: str) -> Any:
|
||||||
|
is_callable, value = self._submit(
|
||||||
|
lambda: (
|
||||||
|
callable(attribute := getattr(self._client, name)),
|
||||||
|
None if callable(attribute) else attribute,
|
||||||
|
)
|
||||||
|
)
|
||||||
|
if not is_callable:
|
||||||
|
return value
|
||||||
|
|
||||||
|
def invoke(*args: Any, **kwargs: Any) -> Any:
|
||||||
|
return self._submit(
|
||||||
|
lambda: getattr(self._client, name)(*args, **kwargs)
|
||||||
|
)
|
||||||
|
|
||||||
|
return invoke
|
||||||
|
|
||||||
|
def close(self) -> None:
|
||||||
|
def close_owned_client() -> None:
|
||||||
|
client, self._client = self._client, None
|
||||||
|
if client is not None:
|
||||||
|
client.close()
|
||||||
|
|
||||||
|
with self._lifecycle_lock:
|
||||||
|
if self._closed:
|
||||||
|
return
|
||||||
|
self._closed = True
|
||||||
|
future = self._executor.submit(close_owned_client)
|
||||||
|
try:
|
||||||
|
future.result()
|
||||||
|
finally:
|
||||||
|
self._executor.shutdown(wait=True, cancel_futures=True)
|
||||||
|
|
||||||
|
|
||||||
class QdrantVectorStore(VectorStoreBase):
|
class QdrantVectorStore(VectorStoreBase):
|
||||||
"""
|
"""
|
||||||
Qdrant 向量存储实现
|
Qdrant 向量存储实现
|
||||||
@@ -113,8 +174,8 @@ class QdrantVectorStore(VectorStoreBase):
|
|||||||
self._is_local = True
|
self._is_local = True
|
||||||
# 兼容性处理: 清理高版本 qdrant-client 写入的 meta.json 中的未知字段
|
# 兼容性处理: 清理高版本 qdrant-client 写入的 meta.json 中的未知字段
|
||||||
self._patch_local_meta(vs_config.persist_directory)
|
self._patch_local_meta(vs_config.persist_directory)
|
||||||
self._client = QdrantClient(
|
self._client = _LocalQdrantOwner(
|
||||||
path=vs_config.persist_directory,
|
lambda: QdrantClient(path=vs_config.persist_directory)
|
||||||
)
|
)
|
||||||
|
|
||||||
# 确保集合存在(同步调用,在 to_thread 中执行)
|
# 确保集合存在(同步调用,在 to_thread 中执行)
|
||||||
@@ -217,7 +278,13 @@ class QdrantVectorStore(VectorStoreBase):
|
|||||||
has_sparse = False
|
has_sparse = False
|
||||||
return has_sparse
|
return has_sparse
|
||||||
|
|
||||||
has_sparse = await _run_in_vdb_pool(_ensure_collection)
|
try:
|
||||||
|
has_sparse = await _run_in_vdb_pool(_ensure_collection)
|
||||||
|
except BaseException:
|
||||||
|
client, self._client = self._client, None
|
||||||
|
if client is not None:
|
||||||
|
await _run_in_vdb_pool(client.close)
|
||||||
|
raise
|
||||||
self._supports_fulltext = bool(has_sparse and bm25_fastembed.is_available())
|
self._supports_fulltext = bool(has_sparse and bm25_fastembed.is_available())
|
||||||
logger.info(
|
logger.info(
|
||||||
f"[qdrant] fulltext(BM25 sparse)={'on' if self._supports_fulltext else 'off'}"
|
f"[qdrant] fulltext(BM25 sparse)={'on' if self._supports_fulltext else 'off'}"
|
||||||
@@ -1043,7 +1110,10 @@ class QdrantVectorStore(VectorStoreBase):
|
|||||||
FieldCondition(key="status", match=MatchAny(any=status_values))
|
FieldCondition(key="status", match=MatchAny(any=status_values))
|
||||||
)
|
)
|
||||||
if layers:
|
if layers:
|
||||||
layer_values = [l.value if hasattr(l, 'value') else str(l) for l in layers]
|
layer_values = [
|
||||||
|
layer.value if hasattr(layer, "value") else str(layer)
|
||||||
|
for layer in layers
|
||||||
|
]
|
||||||
if not allow_raw_layer:
|
if not allow_raw_layer:
|
||||||
layer_values = [v for v in layer_values if v != MemoryLayer.L1_RAW.value]
|
layer_values = [v for v in layer_values if v != MemoryLayer.L1_RAW.value]
|
||||||
if layer_values:
|
if layer_values:
|
||||||
@@ -1372,9 +1442,9 @@ class QdrantVectorStore(VectorStoreBase):
|
|||||||
|
|
||||||
async def close(self) -> None:
|
async def close(self) -> None:
|
||||||
"""关闭客户端"""
|
"""关闭客户端"""
|
||||||
if self._client:
|
client, self._client = self._client, None
|
||||||
self._client.close()
|
if client is not None:
|
||||||
self._client = None
|
await _run_in_vdb_pool(client.close)
|
||||||
|
|
||||||
# ================================================================
|
# ================================================================
|
||||||
# Keyword Search (Full-text)
|
# Keyword Search (Full-text)
|
||||||
@@ -1412,8 +1482,6 @@ class QdrantVectorStore(VectorStoreBase):
|
|||||||
except Exception as e:
|
except Exception as e:
|
||||||
logger.debug(f"[qdrant] inspect search_text index failed (will try create anyway): {e}")
|
logger.debug(f"[qdrant] inspect search_text index failed (will try create anyway): {e}")
|
||||||
|
|
||||||
target_tokenizer_value = getattr(TokenizerType.WHITESPACE, "value", "whitespace")
|
|
||||||
|
|
||||||
if existing_tokenizer is not None:
|
if existing_tokenizer is not None:
|
||||||
if str(existing_tokenizer).lower().endswith("whitespace") or str(existing_tokenizer).lower() == "whitespace":
|
if str(existing_tokenizer).lower().endswith("whitespace") or str(existing_tokenizer).lower() == "whitespace":
|
||||||
return # 已经是新格式,啥都不做
|
return # 已经是新格式,啥都不做
|
||||||
@@ -1436,7 +1504,9 @@ class QdrantVectorStore(VectorStoreBase):
|
|||||||
max_token_len=20,
|
max_token_len=20,
|
||||||
),
|
),
|
||||||
)
|
)
|
||||||
logger.info(f"[qdrant] Created search_text text index (whitespace) on existing collection")
|
logger.info(
|
||||||
|
"[qdrant] Created search_text text index (whitespace) on existing collection"
|
||||||
|
)
|
||||||
except Exception as e:
|
except Exception as e:
|
||||||
# 兜底:未支持或并发竞争,静默
|
# 兜底:未支持或并发竞争,静默
|
||||||
logger.debug(f"[qdrant] _ensure_text_index swallowed: {e}")
|
logger.debug(f"[qdrant] _ensure_text_index swallowed: {e}")
|
||||||
@@ -1510,7 +1580,10 @@ class QdrantVectorStore(VectorStoreBase):
|
|||||||
)
|
)
|
||||||
if layers:
|
if layers:
|
||||||
must_conditions.append(
|
must_conditions.append(
|
||||||
FieldCondition(key="layer", match=MatchAny(any=[l.value for l in layers]))
|
FieldCondition(
|
||||||
|
key="layer",
|
||||||
|
match=MatchAny(any=[layer.value for layer in layers]),
|
||||||
|
)
|
||||||
)
|
)
|
||||||
if status_filter:
|
if status_filter:
|
||||||
must_conditions.append(
|
must_conditions.append(
|
||||||
@@ -1729,7 +1802,7 @@ class QdrantVectorStore(VectorStoreBase):
|
|||||||
|
|
||||||
if changed:
|
if changed:
|
||||||
meta_path.write_text(_json.dumps(meta), encoding="utf-8")
|
meta_path.write_text(_json.dumps(meta), encoding="utf-8")
|
||||||
logger.info(f"Patched Qdrant meta.json: removed unsupported fields")
|
logger.info("Patched Qdrant meta.json: removed unsupported fields")
|
||||||
|
|
||||||
# ================================================================
|
# ================================================================
|
||||||
# Tag Index(per-user tag embedding,供 reader_hybrid_tag 使用)
|
# Tag Index(per-user tag embedding,供 reader_hybrid_tag 使用)
|
||||||
|
|||||||
@@ -0,0 +1,83 @@
|
|||||||
|
"""Local Qdrant must keep its SQLite connection on one owner thread."""
|
||||||
|
|
||||||
|
import asyncio
|
||||||
|
|
||||||
|
import pytest
|
||||||
|
|
||||||
|
qdrant_client = pytest.importorskip("qdrant_client")
|
||||||
|
|
||||||
|
from qdrant_client import QdrantClient, models # noqa: E402
|
||||||
|
from qdrant_client.local.persistence import CollectionPersistence # noqa: E402
|
||||||
|
|
||||||
|
from hy_memory.config import MemoryConfig, VectorStoreConfig # noqa: E402
|
||||||
|
from hy_memory.data.vector_store import create_vector_store # noqa: E402
|
||||||
|
from hy_memory.models.memory import MemoryLayer, MemoryNode # noqa: E402
|
||||||
|
|
||||||
|
|
||||||
|
def _config(directory, name="thread_probe"):
|
||||||
|
config = MemoryConfig()
|
||||||
|
config.vector_store = VectorStoreConfig(
|
||||||
|
provider="qdrant",
|
||||||
|
persist_directory=str(directory),
|
||||||
|
collection_name=name,
|
||||||
|
embedding_dims=4,
|
||||||
|
)
|
||||||
|
return config
|
||||||
|
|
||||||
|
|
||||||
|
def _node(index=0):
|
||||||
|
return MemoryNode(
|
||||||
|
node_id=f"test-node-{index}",
|
||||||
|
user_id="test-user",
|
||||||
|
agent_id="test-agent",
|
||||||
|
session_id="test-session",
|
||||||
|
layer=MemoryLayer.L2_FACT,
|
||||||
|
content=f"Offline owner-thread test {index}",
|
||||||
|
embedding=[1.0, 0.0, 0.0, float(index)],
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.fixture
|
||||||
|
def same_thread_qdrant(monkeypatch):
|
||||||
|
monkeypatch.setattr(CollectionPersistence, "CHECK_SAME_THREAD", True)
|
||||||
|
monkeypatch.setattr(
|
||||||
|
"hy_memory.pipelines._retrieval.bm25_fastembed.is_available",
|
||||||
|
lambda: False,
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_reopened_local_collection_writes_with_same_thread_check(
|
||||||
|
tmp_path,
|
||||||
|
same_thread_qdrant,
|
||||||
|
):
|
||||||
|
client = QdrantClient(path=str(tmp_path))
|
||||||
|
client.create_collection(
|
||||||
|
"thread_probe_4",
|
||||||
|
vectors_config=models.VectorParams(
|
||||||
|
size=4,
|
||||||
|
distance=models.Distance.COSINE,
|
||||||
|
),
|
||||||
|
)
|
||||||
|
client.close()
|
||||||
|
|
||||||
|
store = create_vector_store(_config(tmp_path))
|
||||||
|
await store.initialize()
|
||||||
|
try:
|
||||||
|
assert await store.upsert(_node()) == "test-node-0"
|
||||||
|
finally:
|
||||||
|
await store.close()
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_fresh_local_collection_serializes_concurrent_writes_and_close(
|
||||||
|
tmp_path,
|
||||||
|
same_thread_qdrant,
|
||||||
|
):
|
||||||
|
store = create_vector_store(_config(tmp_path))
|
||||||
|
await store.initialize()
|
||||||
|
try:
|
||||||
|
ids = await asyncio.gather(*(store.upsert(_node(index)) for index in range(8)))
|
||||||
|
assert ids == [f"test-node-{index}" for index in range(8)]
|
||||||
|
finally:
|
||||||
|
await store.close()
|
||||||
Reference in new issue
Block a user