openvector_dev 0.1__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.
- openvector_dev-0.1/PKG-INFO +96 -0
- openvector_dev-0.1/README.md +79 -0
- openvector_dev-0.1/pyproject.toml +29 -0
- openvector_dev-0.1/src/lein_vector/__init__.py +11 -0
- openvector_dev-0.1/src/lein_vector/api/__init__.py +1 -0
- openvector_dev-0.1/src/lein_vector/api/facade.py +136 -0
- openvector_dev-0.1/src/lein_vector/bases/__init__.py +0 -0
- openvector_dev-0.1/src/lein_vector/bases/embeding_provider_abc.py +11 -0
- openvector_dev-0.1/src/lein_vector/bases/memory_manager_abc.py +33 -0
- openvector_dev-0.1/src/lein_vector/memory_manager_qdrant.py +100 -0
- openvector_dev-0.1/src/lein_vector/memory_manager_ram.py +141 -0
- openvector_dev-0.1/src/lein_vector/qdrant_adapter.py +98 -0
- openvector_dev-0.1/src/lein_vector/schemas/__init__.py +0 -0
- openvector_dev-0.1/src/lein_vector/schemas/chunk.py +39 -0
- openvector_dev-0.1/src/lein_vector/sentence_transformer.py +68 -0
- openvector_dev-0.1/src/lein_vector/short_term.py +44 -0
@@ -0,0 +1,96 @@
|
|
1
|
+
Metadata-Version: 2.3
|
2
|
+
Name: openvector_dev
|
3
|
+
Version: 0.1
|
4
|
+
Summary:
|
5
|
+
Author: p00ler
|
6
|
+
Author-email: liveitspain@gmail.com
|
7
|
+
Requires-Python: >=3.11
|
8
|
+
Classifier: Programming Language :: Python :: 3
|
9
|
+
Classifier: Programming Language :: Python :: 3.11
|
10
|
+
Classifier: Programming Language :: Python :: 3.12
|
11
|
+
Classifier: Programming Language :: Python :: 3.13
|
12
|
+
Requires-Dist: aiohttp (>=3.12.13,<4.0.0)
|
13
|
+
Requires-Dist: google-genai[aiohttp] (>=1.21.1,<2.0.0)
|
14
|
+
Requires-Dist: qdrant-client (>=1.14.3,<2.0.0)
|
15
|
+
Description-Content-Type: text/markdown
|
16
|
+
|
17
|
+
# Persona-Memory Subsystem
|
18
|
+
|
19
|
+
## Назначение:
|
20
|
+
Модуль памяти для диалогового ИИ-агента.
|
21
|
+
|
22
|
+
Short-term: последние сообщения (RAM, deque).
|
23
|
+
|
24
|
+
Long-term: Qdrant + архив (долговременные факты и истории).
|
25
|
+
|
26
|
+
## Архитектура:
|
27
|
+
ShortTermMemory — хранит последние сообщения.
|
28
|
+
|
29
|
+
QdrantAdapter — async-слой для поиска/хранения чанков в Qdrant.
|
30
|
+
|
31
|
+
MemoryManagerQdrant — бизнес-логика.
|
32
|
+
|
33
|
+
EmbeddingProviderSentenceTransformer / Gemini — эмбеддинги.
|
34
|
+
|
35
|
+
MemoryService (фасад) — единая точка для верхнего слоя.
|
36
|
+
|
37
|
+
Chunk / ChunkPayload — pydantic-модели для хранения.
|
38
|
+
|
39
|
+
## Пример рабочего цикла:
|
40
|
+
```
|
41
|
+
mem = MemoryService(short_term, memory_manager)
|
42
|
+
mem.add_short("user", user_msg)
|
43
|
+
emb = await embedder.get_embedding(user_msg)
|
44
|
+
long_memories = await mem.get_long(user_id, emb, k=3)
|
45
|
+
short_ctx = mem.get_short(10)
|
46
|
+
mem.add_short("gf", answer)
|
47
|
+
await mem.save_long(user_id, Chunk(
|
48
|
+
chunk_id=uuid4(), user_id=user_id, chunk_type="type0",
|
49
|
+
created_at=datetime.utcnow(), last_hit=datetime.utcnow(),
|
50
|
+
hit_count=0, text=answer, persistent=False
|
51
|
+
))
|
52
|
+
```
|
53
|
+
|
54
|
+
## Установка и тесты:
|
55
|
+
```
|
56
|
+
poetry add ./vector-memory
|
57
|
+
|
58
|
+
docker run -d --name qdrant -p 6333:6333 -v qdrant_data:/qdrant/storage qdrant/qdrant
|
59
|
+
|
60
|
+
pytest
|
61
|
+
```
|
62
|
+
|
63
|
+
## Конфигурация:
|
64
|
+
```
|
65
|
+
QDRANT_HOST, QDRANT_PORT
|
66
|
+
QDRANT_COLLECTION
|
67
|
+
VECTOR_SIZE (совпадает с embedding-моделью)
|
68
|
+
GEMINI_API_KEY (Gemini Embedding-004)
|
69
|
+
```
|
70
|
+
|
71
|
+
## Перед релизом:
|
72
|
+
- Все тесты проходят (pytest)
|
73
|
+
- Размеры векторов и коллекции совпадают
|
74
|
+
- Архивирование/restore, merge, фильтры работают
|
75
|
+
- Нет deprecated-методов в adapter
|
76
|
+
|
77
|
+
## Roadmap:
|
78
|
+
- Redis-кеш для hit_count
|
79
|
+
- Курсорный scroll для больших архивов
|
80
|
+
- gRPC/gateway-адаптер
|
81
|
+
|
82
|
+
## Test-Coverage
|
83
|
+
```
|
84
|
+
Name Stmts Miss Cover
|
85
|
+
-----------------------------------------------------
|
86
|
+
src\bases\memory_manager_abc.py 20 1 95%
|
87
|
+
src\memory_manager_qdrant.py 54 11 80%
|
88
|
+
src\memory_manager_ram.py 95 32 66%
|
89
|
+
src\qdrant_adapter.py 46 8 83%
|
90
|
+
src\schemas\chunk.py 31 1 97%
|
91
|
+
src\short_term.py 26 1 96%
|
92
|
+
-----------------------------------------------------
|
93
|
+
TOTAL 272 54 80%
|
94
|
+
```
|
95
|
+
|
96
|
+
|
@@ -0,0 +1,79 @@
|
|
1
|
+
# Persona-Memory Subsystem
|
2
|
+
|
3
|
+
## Назначение:
|
4
|
+
Модуль памяти для диалогового ИИ-агента.
|
5
|
+
|
6
|
+
Short-term: последние сообщения (RAM, deque).
|
7
|
+
|
8
|
+
Long-term: Qdrant + архив (долговременные факты и истории).
|
9
|
+
|
10
|
+
## Архитектура:
|
11
|
+
ShortTermMemory — хранит последние сообщения.
|
12
|
+
|
13
|
+
QdrantAdapter — async-слой для поиска/хранения чанков в Qdrant.
|
14
|
+
|
15
|
+
MemoryManagerQdrant — бизнес-логика.
|
16
|
+
|
17
|
+
EmbeddingProviderSentenceTransformer / Gemini — эмбеддинги.
|
18
|
+
|
19
|
+
MemoryService (фасад) — единая точка для верхнего слоя.
|
20
|
+
|
21
|
+
Chunk / ChunkPayload — pydantic-модели для хранения.
|
22
|
+
|
23
|
+
## Пример рабочего цикла:
|
24
|
+
```
|
25
|
+
mem = MemoryService(short_term, memory_manager)
|
26
|
+
mem.add_short("user", user_msg)
|
27
|
+
emb = await embedder.get_embedding(user_msg)
|
28
|
+
long_memories = await mem.get_long(user_id, emb, k=3)
|
29
|
+
short_ctx = mem.get_short(10)
|
30
|
+
mem.add_short("gf", answer)
|
31
|
+
await mem.save_long(user_id, Chunk(
|
32
|
+
chunk_id=uuid4(), user_id=user_id, chunk_type="type0",
|
33
|
+
created_at=datetime.utcnow(), last_hit=datetime.utcnow(),
|
34
|
+
hit_count=0, text=answer, persistent=False
|
35
|
+
))
|
36
|
+
```
|
37
|
+
|
38
|
+
## Установка и тесты:
|
39
|
+
```
|
40
|
+
poetry add ./vector-memory
|
41
|
+
|
42
|
+
docker run -d --name qdrant -p 6333:6333 -v qdrant_data:/qdrant/storage qdrant/qdrant
|
43
|
+
|
44
|
+
pytest
|
45
|
+
```
|
46
|
+
|
47
|
+
## Конфигурация:
|
48
|
+
```
|
49
|
+
QDRANT_HOST, QDRANT_PORT
|
50
|
+
QDRANT_COLLECTION
|
51
|
+
VECTOR_SIZE (совпадает с embedding-моделью)
|
52
|
+
GEMINI_API_KEY (Gemini Embedding-004)
|
53
|
+
```
|
54
|
+
|
55
|
+
## Перед релизом:
|
56
|
+
- Все тесты проходят (pytest)
|
57
|
+
- Размеры векторов и коллекции совпадают
|
58
|
+
- Архивирование/restore, merge, фильтры работают
|
59
|
+
- Нет deprecated-методов в adapter
|
60
|
+
|
61
|
+
## Roadmap:
|
62
|
+
- Redis-кеш для hit_count
|
63
|
+
- Курсорный scroll для больших архивов
|
64
|
+
- gRPC/gateway-адаптер
|
65
|
+
|
66
|
+
## Test-Coverage
|
67
|
+
```
|
68
|
+
Name Stmts Miss Cover
|
69
|
+
-----------------------------------------------------
|
70
|
+
src\bases\memory_manager_abc.py 20 1 95%
|
71
|
+
src\memory_manager_qdrant.py 54 11 80%
|
72
|
+
src\memory_manager_ram.py 95 32 66%
|
73
|
+
src\qdrant_adapter.py 46 8 83%
|
74
|
+
src\schemas\chunk.py 31 1 97%
|
75
|
+
src\short_term.py 26 1 96%
|
76
|
+
-----------------------------------------------------
|
77
|
+
TOTAL 272 54 80%
|
78
|
+
```
|
79
|
+
|
@@ -0,0 +1,29 @@
|
|
1
|
+
[project]
|
2
|
+
name = "openvector_dev"
|
3
|
+
version = "0.1"
|
4
|
+
description = ""
|
5
|
+
authors = [
|
6
|
+
{name = "p00ler",email = "liveitspain@gmail.com"}
|
7
|
+
]
|
8
|
+
readme = "README.md"
|
9
|
+
requires-python = ">=3.11"
|
10
|
+
dependencies = [
|
11
|
+
"qdrant-client (>=1.14.3,<2.0.0)",
|
12
|
+
"aiohttp (>=3.12.13,<4.0.0)",
|
13
|
+
"google-genai[aiohttp] (>=1.21.1,<2.0.0)",
|
14
|
+
]
|
15
|
+
|
16
|
+
|
17
|
+
[build-system]
|
18
|
+
requires = ["poetry-core>=2.0.0,<3.0.0"]
|
19
|
+
build-backend = "poetry.core.masonry.api"
|
20
|
+
|
21
|
+
[tool.poetry.group.dev.dependencies]
|
22
|
+
pytest = "^8.4.1"
|
23
|
+
pytest-asyncio = "^1.0.0"
|
24
|
+
pytest-cov = "^6.2.1"
|
25
|
+
|
26
|
+
[tool.poetry]
|
27
|
+
packages = [
|
28
|
+
{ include = "lein_vector", from = "src" }
|
29
|
+
]
|
@@ -0,0 +1,11 @@
|
|
1
|
+
from . import api, bases, schemas
|
2
|
+
from .memory_manager_qdrant import MemoryManagerQdrant
|
3
|
+
from .qdrant_adapter import QdrantAdapter
|
4
|
+
from .sentence_transformer import EmbeddingProviderGemini as EmbeddingProvider
|
5
|
+
from .short_term import ShortTermMemory
|
6
|
+
|
7
|
+
__all__ = [
|
8
|
+
"api", "bases", "schemas",
|
9
|
+
"MemoryManagerQdrant",
|
10
|
+
"QdrantAdapter", "EmbeddingProvider", "ShortTermMemory"
|
11
|
+
]
|
@@ -0,0 +1 @@
|
|
1
|
+
from facade import MemoryFacade as Memory
|
@@ -0,0 +1,136 @@
|
|
1
|
+
from uuid import uuid4
|
2
|
+
from datetime import datetime, UTC
|
3
|
+
|
4
|
+
from src.lein_vector import ShortTermMemory, QdrantAdapter, MemoryManagerQdrant
|
5
|
+
from src.lein_vector.schemas.chunk import Chunk
|
6
|
+
from src.lein_vector.sentence_transformer import EmbeddingProviderGemini
|
7
|
+
|
8
|
+
|
9
|
+
class MemoryFacade:
|
10
|
+
def __init__(self, short_term, memory_manager, embedder):
|
11
|
+
self.short = short_term
|
12
|
+
self.long = memory_manager
|
13
|
+
self.embed = embedder
|
14
|
+
self._msg_no: dict[int, int] = {}
|
15
|
+
|
16
|
+
@classmethod
|
17
|
+
async def from_qdrant(
|
18
|
+
cls,
|
19
|
+
host: str,
|
20
|
+
port: int,
|
21
|
+
collection: str,
|
22
|
+
vector_size: int = 768,
|
23
|
+
api_key: str | None = None,
|
24
|
+
short_maxlen: int = 20,
|
25
|
+
) -> "MemoryFacade":
|
26
|
+
"""
|
27
|
+
Создаёт MemoryFacade со всеми зависимостями:
|
28
|
+
- ShortTermMemory(maxlen=short_maxlen)
|
29
|
+
- EmbeddingProviderGemini(api_key)
|
30
|
+
- QdrantAdapter(host, port, collection, vector_size) + init_collection()
|
31
|
+
- MemoryManagerQdrant(adapter, embedder)
|
32
|
+
"""
|
33
|
+
# 1. short-term
|
34
|
+
short_mem = ShortTermMemory(maxlen=short_maxlen)
|
35
|
+
|
36
|
+
# 2. эмбеддер
|
37
|
+
embedder = EmbeddingProviderGemini(api_key=api_key)
|
38
|
+
|
39
|
+
# 3. адаптер Qdrant
|
40
|
+
adapter = QdrantAdapter(host, port, collection, vector_size)
|
41
|
+
await adapter.init_collection()
|
42
|
+
|
43
|
+
# 4. менеджер долгой памяти
|
44
|
+
long_mem = MemoryManagerQdrant(adapter, embedder)
|
45
|
+
|
46
|
+
# 5. возвращаем фасад
|
47
|
+
return cls(short_mem, long_mem, embedder)
|
48
|
+
|
49
|
+
async def step_user(self, user_id: int, user_msg: str, topk: int = 3, history_n: int = 20):
|
50
|
+
self.short.add("user", user_msg)
|
51
|
+
embedding = await self.embed.get_embedding(user_msg)
|
52
|
+
long_memories = await self.long.retrieve_by_embedding(user_id, embedding, topk=topk)
|
53
|
+
short_ctx = self.short.window(history_n)
|
54
|
+
return {
|
55
|
+
"short_term": short_ctx,
|
56
|
+
"long_term": long_memories
|
57
|
+
}
|
58
|
+
|
59
|
+
async def step_user_oai(
|
60
|
+
self,
|
61
|
+
user_id: int,
|
62
|
+
user_msg: str,
|
63
|
+
*,
|
64
|
+
topk: int = 3,
|
65
|
+
history_n: int = 20,
|
66
|
+
) -> dict:
|
67
|
+
"""
|
68
|
+
Полный шаг для OpenAI-совместимого вывода:
|
69
|
+
1. Записывает сообщение пользователя в short-term.
|
70
|
+
2. Достаёт релевантные чанки из long-term.
|
71
|
+
3. Возвращает short-term уже в формате OpenAI.
|
72
|
+
"""
|
73
|
+
data = await self.step_user(user_id, user_msg, topk=topk, history_n=history_n)
|
74
|
+
data["short_term"] = self._to_openai(data["short_term"])
|
75
|
+
return data
|
76
|
+
|
77
|
+
@staticmethod
|
78
|
+
def _to_openai(msgs: list[dict]) -> list[dict]:
|
79
|
+
role_map = {"gf": "assistant"} # «gf» → OpenAI «assistant»
|
80
|
+
return [
|
81
|
+
{"role": role_map.get(m["role"], m["role"]), "content": m["text"]}
|
82
|
+
for m in msgs
|
83
|
+
]
|
84
|
+
|
85
|
+
async def step_gf(
|
86
|
+
self,
|
87
|
+
user_id: int,
|
88
|
+
gf_msg: str,
|
89
|
+
*,
|
90
|
+
block_size: int = 8,
|
91
|
+
save_pair: bool = True,
|
92
|
+
):
|
93
|
+
# 1) кладём ответ в short-term
|
94
|
+
curr_no = self._msg_no.get(user_id, 0) + 1
|
95
|
+
self._msg_no[user_id] = curr_no
|
96
|
+
self.short.add("gf", gf_msg, extra={"msg_no": curr_no})
|
97
|
+
|
98
|
+
# 2) если блок из 'block_size' сообщений готов → формируем long-term чанк
|
99
|
+
if save_pair and len(self.short.window()) >= block_size:
|
100
|
+
last_block = self.short.window(block_size) # последние 8 сообщений
|
101
|
+
block_text = "\n".join(m["text"] for m in last_block)
|
102
|
+
|
103
|
+
# считаем embedding один раз
|
104
|
+
vector = await self.embed.get_embedding(block_text)
|
105
|
+
|
106
|
+
new_chunk = Chunk(
|
107
|
+
chunk_id=uuid4(),
|
108
|
+
user_id=user_id,
|
109
|
+
chunk_type="type0",
|
110
|
+
created_at=datetime.now(UTC),
|
111
|
+
last_hit=datetime.now(UTC),
|
112
|
+
hit_count=0,
|
113
|
+
text=block_text,
|
114
|
+
persistent=False,
|
115
|
+
extra={"msg_no": curr_no},
|
116
|
+
)
|
117
|
+
await self.long.upsert_chunk_with_vector(user_id, new_chunk, vector)
|
118
|
+
|
119
|
+
# (необязательно) можешь очистить short-term, если maxlen маленький
|
120
|
+
# self.short.clear_until(block_size) ← если нужен скользящий сдвиг
|
121
|
+
|
122
|
+
# 3) при необходимости запускаем merge / maintenance
|
123
|
+
if curr_no % 40 == 0: # каждые 40 сообщений
|
124
|
+
await self.long.merge_old_chunks(user_id, "type0", n=5)
|
125
|
+
|
126
|
+
def get_short_term(self, n=10) -> list:
|
127
|
+
return self.short.window(n)
|
128
|
+
|
129
|
+
async def get_long_term(self, user_id: int, embedding: list[float], topk: int = 3) -> list:
|
130
|
+
return await self.long.retrieve_by_embedding(user_id, embedding, topk)
|
131
|
+
|
132
|
+
def add_to_short(self, role: str, text: str) -> None:
|
133
|
+
self.short.add(role, text)
|
134
|
+
|
135
|
+
async def add_to_long(self, user_id: int, chunk: Chunk) -> None:
|
136
|
+
await self.long.upsert_chunk(user_id, chunk)
|
File without changes
|
@@ -0,0 +1,11 @@
|
|
1
|
+
from typing import List
|
2
|
+
from abc import ABC, abstractmethod
|
3
|
+
|
4
|
+
class EmbeddingProviderABC(ABC):
|
5
|
+
@abstractmethod
|
6
|
+
async def get_embedding(self, text: str) -> List[float]:
|
7
|
+
...
|
8
|
+
|
9
|
+
@abstractmethod
|
10
|
+
async def get_embeddings(self, texts: List[str]) -> List[List[float]]:
|
11
|
+
...
|
@@ -0,0 +1,33 @@
|
|
1
|
+
from typing import List
|
2
|
+
from uuid import UUID
|
3
|
+
from src.lein_vector.schemas.chunk import Chunk
|
4
|
+
|
5
|
+
|
6
|
+
class MemoryManagerABC:
|
7
|
+
async def upsert_chunk(self, user_id: int, chunk: Chunk) -> None: ...
|
8
|
+
async def upsert_chunks(self, user_id: int, chunks: List[Chunk]) -> None: ...
|
9
|
+
async def retrieve_by_embedding(
|
10
|
+
self, user_id: int, embedding: List[float], topk: int = 3
|
11
|
+
) -> List[Chunk]: ...
|
12
|
+
async def retrieve_by_embedding_batch(
|
13
|
+
self, user_id: int, embeddings: List[List[float]], topk: int = 3
|
14
|
+
) -> List[List[Chunk]]:
|
15
|
+
raise NotImplementedError("Not implemented in this backend")
|
16
|
+
|
17
|
+
async def retrieve_by_type(
|
18
|
+
self, user_id: int, chunk_type: str, topk: int = 3
|
19
|
+
) -> List[Chunk]: ...
|
20
|
+
async def retrieve_by_text(
|
21
|
+
self, user_id: int, query: str, topk: int = 3
|
22
|
+
) -> List[Chunk]: ...
|
23
|
+
async def merge_old_chunks(self, user_id: int, chunk_type: str) -> None: ...
|
24
|
+
async def archive_user(self, user_id: int) -> None: ...
|
25
|
+
async def restore_user(self, user_id: int) -> None: ...
|
26
|
+
async def increment_hit(self, user_id: int, chunk_id: UUID) -> None: ...
|
27
|
+
async def pop_first_n(
|
28
|
+
self, user_id: int, chunk_type: str, n: int
|
29
|
+
) -> List[Chunk]: ...
|
30
|
+
async def delete_oldest_nonpersistent(self, user_id: int, keep: int) -> None: ...
|
31
|
+
async def delete_chunk(self, user_id: int, chunk_id: UUID) -> None: ...
|
32
|
+
async def delete_chunks(self, user_id: int, chunk_ids: List[UUID]) -> None: ...
|
33
|
+
async def delete_all(self, user_id: int) -> None: ...
|
@@ -0,0 +1,100 @@
|
|
1
|
+
from datetime import datetime, UTC
|
2
|
+
from typing import List, Dict, Any
|
3
|
+
from uuid import UUID
|
4
|
+
|
5
|
+
from src.lein_vector.bases.memory_manager_abc import MemoryManagerABC
|
6
|
+
from src.lein_vector.schemas.chunk import Chunk, ChunkPayload
|
7
|
+
|
8
|
+
class MemoryManagerQdrant(MemoryManagerABC):
|
9
|
+
def __init__(self, qdrant_adapter, embedding_provider, archive_storage=None):
|
10
|
+
self.qdrant = qdrant_adapter
|
11
|
+
self.embed = embedding_provider
|
12
|
+
self.archive = archive_storage # твой модуль S3/minio (интерфейс: save(user_id, List[ChunkPayload]), load(user_id) -> List[ChunkPayload])
|
13
|
+
|
14
|
+
async def upsert_chunk(self, user_id: int, chunk: Chunk) -> None:
|
15
|
+
embedding = await self.embed.get_embedding(chunk.text)
|
16
|
+
await self.qdrant.upsert(chunk.chunk_id, embedding, chunk.to_payload())
|
17
|
+
|
18
|
+
async def upsert_chunk_with_vector(
|
19
|
+
self,
|
20
|
+
user_id: int,
|
21
|
+
chunk: Chunk,
|
22
|
+
embedding: list[float]
|
23
|
+
) -> None:
|
24
|
+
await self.qdrant.upsert(chunk.chunk_id, embedding, chunk.to_payload())
|
25
|
+
|
26
|
+
async def upsert_chunks(self, user_id: int, chunks: List[Chunk]) -> None:
|
27
|
+
texts = [c.text for c in chunks]
|
28
|
+
embeddings = await self.embed.get_embeddings(texts)
|
29
|
+
points = [
|
30
|
+
{"point_id": c.chunk_id, "embedding": emb, "payload": c.to_payload()}
|
31
|
+
for c, emb in zip(chunks, embeddings)
|
32
|
+
]
|
33
|
+
await self.qdrant.upsert_batch(points)
|
34
|
+
|
35
|
+
async def retrieve_by_embedding(self, user_id: int, embedding: List[float], topk: int = 3, filter_: Dict[str, Any] = None) -> List[ChunkPayload]:
|
36
|
+
# Фильтр по user_id + кастомные условия
|
37
|
+
filter_ = filter_ or {}
|
38
|
+
filter_["user_id"] = user_id
|
39
|
+
return await self.qdrant.search(embedding, filter_, topk)
|
40
|
+
|
41
|
+
async def retrieve_by_type(self, user_id: int, chunk_type: str, topk: int = 3) -> List[ChunkPayload]:
|
42
|
+
# Заглушка embedding (пустой вектор не сработает, нужно реальный запрос!):
|
43
|
+
# Лучше использовать scroll по фильтру
|
44
|
+
filter_ = {"user_id": user_id, "chunk_type": chunk_type}
|
45
|
+
return await self.qdrant.get_all_chunks_with_filter(filter_)
|
46
|
+
|
47
|
+
async def merge_old_chunks(self, user_id: int, chunk_type: str, n: int = 5) -> None:
|
48
|
+
# 1. Получить n старых чанков нужного типа
|
49
|
+
chunks = await self.qdrant.get_n_oldest_chunks(user_id, chunk_type, n)
|
50
|
+
if len(chunks) < n:
|
51
|
+
return
|
52
|
+
# 2. Суммаризация (mock или через LLM)
|
53
|
+
merged_text = " | ".join([c.text for c in chunks])
|
54
|
+
from uuid import uuid4
|
55
|
+
from datetime import datetime
|
56
|
+
summary_chunk = Chunk(
|
57
|
+
chunk_id=uuid4(),
|
58
|
+
user_id=user_id,
|
59
|
+
chunk_type=self._next_type(chunk_type),
|
60
|
+
created_at=datetime.now(UTC),
|
61
|
+
last_hit=datetime.now(UTC),
|
62
|
+
hit_count=0,
|
63
|
+
text=merged_text,
|
64
|
+
persistent=False,
|
65
|
+
summary_of=[c.chunk_id for c in chunks],
|
66
|
+
)
|
67
|
+
await self.upsert_chunk(user_id, summary_chunk)
|
68
|
+
# 3. Удалить исходники
|
69
|
+
await self.delete_chunks(user_id, [c.chunk_id for c in chunks])
|
70
|
+
|
71
|
+
async def archive_user(self, user_id: int) -> None:
|
72
|
+
all_chunks = await self.qdrant.get_all_chunks(user_id)
|
73
|
+
await self.archive.save(user_id, all_chunks)
|
74
|
+
await self.delete_all(user_id)
|
75
|
+
|
76
|
+
async def restore_user(self, user_id: int) -> None:
|
77
|
+
chunks = await self.archive.load(user_id)
|
78
|
+
await self.upsert_chunks(
|
79
|
+
user_id,
|
80
|
+
[Chunk(**c.dict(), last_hit=datetime.now(UTC), hit_count=0) for c in chunks]
|
81
|
+
)
|
82
|
+
|
83
|
+
async def delete_chunk(self, user_id: int, chunk_id: UUID) -> None:
|
84
|
+
await self.qdrant.delete(chunk_id)
|
85
|
+
|
86
|
+
async def delete_chunks(self, user_id: int, chunk_ids: List[UUID]) -> None:
|
87
|
+
await self.qdrant.delete_batch(chunk_ids)
|
88
|
+
|
89
|
+
async def delete_all(self, user_id: int) -> None:
|
90
|
+
all_chunks = await self.qdrant.get_all_chunks(user_id)
|
91
|
+
await self.delete_chunks(user_id, [c.chunk_id for c in all_chunks])
|
92
|
+
|
93
|
+
# Доп. методы поиска (по времени, hit_count, last_hit)
|
94
|
+
async def retrieve_filtered(self, user_id: int, filter_: Dict[str, Any], topk: int = 10) -> List[ChunkPayload]:
|
95
|
+
return await self.qdrant.get_all_chunks_with_filter({"user_id": user_id, **filter_}, topk=topk)
|
96
|
+
|
97
|
+
def _next_type(self, chunk_type: str) -> str:
|
98
|
+
# Логика типа next_type
|
99
|
+
mapping = {"type0": "type1", "type1": "type2"}
|
100
|
+
return mapping.get(chunk_type, "summary")
|
@@ -0,0 +1,141 @@
|
|
1
|
+
from typing import List
|
2
|
+
from uuid import UUID
|
3
|
+
from src.lein_vector.schemas.chunk import Chunk
|
4
|
+
from src.lein_vector.bases.memory_manager_abc import MemoryManagerABC
|
5
|
+
from uuid import uuid4
|
6
|
+
from datetime import datetime, timezone
|
7
|
+
|
8
|
+
class MemoryManagerRAM(MemoryManagerABC):
|
9
|
+
def __init__(self):
|
10
|
+
self._data: dict[int, dict[UUID, Chunk]] = {}
|
11
|
+
self._archive: dict[int, dict[UUID, Chunk]] = {}
|
12
|
+
|
13
|
+
async def upsert_chunk(self, user_id: int, chunk: Chunk) -> None:
|
14
|
+
if user_id not in self._data:
|
15
|
+
self._data[user_id] = {}
|
16
|
+
self._data[user_id][chunk.chunk_id] = chunk
|
17
|
+
|
18
|
+
async def upsert_chunks(self, user_id: int, chunks: List[Chunk]) -> None:
|
19
|
+
if user_id not in self._data:
|
20
|
+
self._data[user_id] = {}
|
21
|
+
for chunk in chunks:
|
22
|
+
self._data[user_id][chunk.chunk_id] = chunk
|
23
|
+
|
24
|
+
async def retrieve_by_embedding(self, user_id: int, embedding: List[float], topk: int = 3) -> List[Chunk]:
|
25
|
+
user_chunks = self._data.get(user_id, {})
|
26
|
+
sorted_chunks = sorted(user_chunks.values(), key=lambda c: c.created_at, reverse=True)
|
27
|
+
return sorted_chunks[:topk]
|
28
|
+
|
29
|
+
async def retrieve_by_embedding_batch(self, user_id: int, embeddings: List[List[float]], topk: int = 3) -> List[List[Chunk]]:
|
30
|
+
raise NotImplementedError("Not implemented in RAM backend")
|
31
|
+
|
32
|
+
async def retrieve_by_type(self, user_id: int, chunk_type: str, topk: int = 3) -> List[Chunk]:
|
33
|
+
user_chunks = self._data.get(user_id, {})
|
34
|
+
filtered = [c for c in user_chunks.values() if c.chunk_type == chunk_type]
|
35
|
+
filtered.sort(key=lambda c: c.created_at, reverse=True)
|
36
|
+
return filtered[:topk]
|
37
|
+
|
38
|
+
async def retrieve_by_text(self, user_id: int, query: str, topk: int = 3) -> List[Chunk]:
|
39
|
+
user_chunks = self._data.get(user_id, {})
|
40
|
+
filtered = [c for c in user_chunks.values() if query.lower() in c.text.lower()]
|
41
|
+
filtered.sort(key=lambda c: c.created_at, reverse=True)
|
42
|
+
return filtered[:topk]
|
43
|
+
|
44
|
+
async def merge_old_chunks(self, user_id: int, chunk_type: str, n: int = 5) -> None:
|
45
|
+
user_chunks = self._data.get(user_id, {})
|
46
|
+
next_type = {
|
47
|
+
"type0": "type1",
|
48
|
+
"type1": "type2"
|
49
|
+
}.get(chunk_type)
|
50
|
+
if not next_type:
|
51
|
+
return
|
52
|
+
|
53
|
+
candidates = [c for c in user_chunks.values() if c.chunk_type == chunk_type]
|
54
|
+
if len(candidates) < n:
|
55
|
+
return
|
56
|
+
|
57
|
+
candidates.sort(key=lambda c: c.created_at)
|
58
|
+
selected = candidates[:n]
|
59
|
+
|
60
|
+
merged_text = " | ".join([c.text for c in selected]) # mock summary
|
61
|
+
|
62
|
+
new_chunk = Chunk(
|
63
|
+
chunk_id=uuid4(),
|
64
|
+
user_id=user_id,
|
65
|
+
chunk_type=next_type,
|
66
|
+
created_at=datetime.now(timezone.utc),
|
67
|
+
last_hit=datetime.now(timezone.utc),
|
68
|
+
hit_count=0,
|
69
|
+
text=merged_text,
|
70
|
+
persistent=False,
|
71
|
+
summary_of=[c.chunk_id for c in selected],
|
72
|
+
)
|
73
|
+
for c in selected:
|
74
|
+
del user_chunks[c.chunk_id]
|
75
|
+
user_chunks[new_chunk.chunk_id] = new_chunk
|
76
|
+
|
77
|
+
async def archive_user(self, user_id: int) -> None:
|
78
|
+
if user_id in self._data:
|
79
|
+
self._archive[user_id] = self._data[user_id]
|
80
|
+
del self._data[user_id]
|
81
|
+
|
82
|
+
async def restore_user(self, user_id: int) -> None:
|
83
|
+
if user_id in self._archive:
|
84
|
+
self._data[user_id] = self._archive[user_id]
|
85
|
+
del self._archive[user_id]
|
86
|
+
|
87
|
+
async def increment_hit(self, user_id: int, chunk_id: UUID) -> None:
|
88
|
+
user_chunks = self._data.get(user_id, {})
|
89
|
+
chunk = user_chunks.get(chunk_id)
|
90
|
+
if chunk is not None:
|
91
|
+
chunk.hit_count += 1
|
92
|
+
from datetime import datetime, timezone
|
93
|
+
chunk.last_hit = datetime.now(timezone.utc)
|
94
|
+
|
95
|
+
async def pop_first_n(self, user_id: int, chunk_type: str, n: int) -> List[Chunk]:
|
96
|
+
user_chunks = self._data.get(user_id, {})
|
97
|
+
filtered = [c for c in user_chunks.values() if c.chunk_type == chunk_type]
|
98
|
+
# сортировка по created_at (старые — первые)
|
99
|
+
filtered.sort(key=lambda c: c.created_at)
|
100
|
+
# выбираем первые n
|
101
|
+
selected = filtered[:n]
|
102
|
+
# удаляем их из данных
|
103
|
+
for chunk in selected:
|
104
|
+
del self._data[user_id][chunk.chunk_id]
|
105
|
+
return selected
|
106
|
+
|
107
|
+
async def delete_oldest_nonpersistent(self, user_id: int, keep: int) -> None:
|
108
|
+
user_chunks = self._data.get(user_id, {})
|
109
|
+
nonpersistent = [c for c in user_chunks.values() if not c.persistent]
|
110
|
+
# сортировка по created_at (старые — первые)
|
111
|
+
nonpersistent.sort(key=lambda c: c.created_at)
|
112
|
+
# если их больше чем keep — удаляем лишние
|
113
|
+
for chunk in nonpersistent[:-keep]:
|
114
|
+
del self._data[user_id][chunk.chunk_id]
|
115
|
+
|
116
|
+
|
117
|
+
async def delete_chunk(self, user_id: int, chunk_id: UUID) -> None:
|
118
|
+
user_chunks = self._data.get(user_id, {})
|
119
|
+
user_chunks.pop(chunk_id, None)
|
120
|
+
|
121
|
+
|
122
|
+
async def delete_chunks(self, user_id: int, chunk_ids: List[UUID]) -> None:
|
123
|
+
user_chunks = self._data.get(user_id, {})
|
124
|
+
for chunk_id in chunk_ids:
|
125
|
+
user_chunks.pop(chunk_id, None)
|
126
|
+
|
127
|
+
async def delete_all(self, user_id: int) -> None:
|
128
|
+
self._data.pop(user_id, None)
|
129
|
+
|
130
|
+
def get_all_chunks(self, user_id: int) -> List[Chunk]:
|
131
|
+
"""Для тестов — все чанки пользователя."""
|
132
|
+
return list(self._data.get(user_id, {}).values())
|
133
|
+
|
134
|
+
def get_all_archive(self, user_id: int) -> List[Chunk]:
|
135
|
+
"""Для тестов — все чанки в архиве."""
|
136
|
+
return list(self._archive.get(user_id, {}).values())
|
137
|
+
|
138
|
+
def clear(self):
|
139
|
+
"""Очистка всех данных для тестов."""
|
140
|
+
self._data.clear()
|
141
|
+
self._archive.clear()
|
@@ -0,0 +1,98 @@
|
|
1
|
+
from qdrant_client import AsyncQdrantClient
|
2
|
+
from qdrant_client.http.models import MatchText
|
3
|
+
from qdrant_client.models import (
|
4
|
+
PointStruct, Filter, FieldCondition, MatchValue, Range,
|
5
|
+
VectorParams, Distance
|
6
|
+
)
|
7
|
+
from src.lein_vector.schemas.chunk import ChunkPayload
|
8
|
+
from typing import List, Dict, Any
|
9
|
+
from uuid import UUID
|
10
|
+
|
11
|
+
|
12
|
+
class QdrantAdapter:
|
13
|
+
def __init__(self, host: str, port: int, collection: str = "persona_mem", vector_size: int = 768):
|
14
|
+
self.collection = collection
|
15
|
+
self.client = AsyncQdrantClient(host=host, port=port)
|
16
|
+
self.vector_size = vector_size
|
17
|
+
|
18
|
+
async def init_collection(self):
|
19
|
+
exists = await self.client.collection_exists(self.collection)
|
20
|
+
if not exists:
|
21
|
+
await self.client.create_collection(
|
22
|
+
collection_name=self.collection,
|
23
|
+
vectors_config=VectorParams(size=self.vector_size, distance=Distance.COSINE)
|
24
|
+
)
|
25
|
+
|
26
|
+
async def upsert(self, point_id: UUID, embedding: List[float], payload: ChunkPayload) -> None:
|
27
|
+
await self.client.upsert(
|
28
|
+
collection_name=self.collection,
|
29
|
+
points=[
|
30
|
+
PointStruct(
|
31
|
+
id=str(point_id),
|
32
|
+
vector=embedding,
|
33
|
+
payload=payload.model_dump()
|
34
|
+
)
|
35
|
+
]
|
36
|
+
)
|
37
|
+
|
38
|
+
async def upsert_batch(self, points: List[Dict[str, Any]]) -> None:
|
39
|
+
structs = [
|
40
|
+
PointStruct(
|
41
|
+
id=str(point["point_id"]),
|
42
|
+
vector=point["embedding"],
|
43
|
+
payload=point["payload"].dict()
|
44
|
+
)
|
45
|
+
for point in points
|
46
|
+
]
|
47
|
+
await self.client.upsert(collection_name=self.collection, points=structs)
|
48
|
+
|
49
|
+
async def search(self, embedding: List[float], filter_: Dict[str, Any], topk: int) -> List[ChunkPayload]:
|
50
|
+
# Пример фильтра {"user_id": 123, "chunk_type": "type1", "created_at_gt": "2024-01-01T00:00:00"}
|
51
|
+
conditions = []
|
52
|
+
for k, v in filter_.items():
|
53
|
+
if k.endswith("_gt"):
|
54
|
+
field = k[:-3]
|
55
|
+
conditions.append(FieldCondition(key=field, range=Range(gt=v)))
|
56
|
+
elif k.endswith("_lt"):
|
57
|
+
field = k[:-3]
|
58
|
+
conditions.append(FieldCondition(key=field, range=Range(lt=v)))
|
59
|
+
else:
|
60
|
+
if isinstance(v, str):
|
61
|
+
conditions.append(FieldCondition(key=k, match=MatchText(text=v)))
|
62
|
+
else:
|
63
|
+
conditions.append(FieldCondition(key=k, match=MatchValue(value=v)))
|
64
|
+
q_filter = Filter(must=conditions)
|
65
|
+
result = await self.client.query_points(
|
66
|
+
collection_name=self.collection,
|
67
|
+
query=embedding,
|
68
|
+
query_filter=q_filter,
|
69
|
+
limit=topk,
|
70
|
+
)
|
71
|
+
points = result.points
|
72
|
+
return [ChunkPayload(**point.payload) for point in points]
|
73
|
+
|
74
|
+
async def delete(self, point_id: UUID) -> None:
|
75
|
+
await self.client.delete(
|
76
|
+
collection_name=self.collection,
|
77
|
+
points_selector=[str(point_id)]
|
78
|
+
)
|
79
|
+
|
80
|
+
async def delete_batch(self, point_ids: List[UUID]) -> None:
|
81
|
+
await self.client.delete(
|
82
|
+
collection_name=self.collection,
|
83
|
+
points_selector=[str(pid) for pid in point_ids]
|
84
|
+
)
|
85
|
+
|
86
|
+
async def delete_collection(self) -> None:
|
87
|
+
await self.client.delete_collection(collection_name=self.collection)
|
88
|
+
|
89
|
+
async def get_all_chunks(self, user_id: int) -> List[ChunkPayload]:
|
90
|
+
q_filter = Filter(
|
91
|
+
must=[FieldCondition(key="user_id", match=MatchValue(value=user_id))]
|
92
|
+
)
|
93
|
+
scroll = await self.client.scroll(
|
94
|
+
collection_name=self.collection,
|
95
|
+
scroll_filter=q_filter,
|
96
|
+
limit=2048,
|
97
|
+
)
|
98
|
+
return [ChunkPayload(**p.payload) for p in scroll[0]]
|
File without changes
|
@@ -0,0 +1,39 @@
|
|
1
|
+
from __future__ import annotations
|
2
|
+
|
3
|
+
from pydantic import BaseModel, Field
|
4
|
+
from typing import List, Optional, Dict
|
5
|
+
from datetime import datetime, UTC
|
6
|
+
from uuid import UUID
|
7
|
+
|
8
|
+
class Chunk(BaseModel):
|
9
|
+
chunk_id: UUID
|
10
|
+
user_id: int
|
11
|
+
chunk_type: str # "type0" | "type1" | "fact"
|
12
|
+
created_at: datetime
|
13
|
+
last_hit: datetime
|
14
|
+
hit_count: int = 0
|
15
|
+
text: str
|
16
|
+
persistent: bool = False
|
17
|
+
|
18
|
+
summary_of: Optional[List[UUID]] = None # для type1
|
19
|
+
source_chunk_id: Optional[UUID] = None # для fact
|
20
|
+
extra: Optional[Dict] = Field(default_factory=dict)
|
21
|
+
|
22
|
+
def to_payload(self) -> ChunkPayload:
|
23
|
+
return ChunkPayload(**self.model_dump())
|
24
|
+
|
25
|
+
|
26
|
+
class ChunkPayload(BaseModel):
|
27
|
+
chunk_id: UUID
|
28
|
+
user_id: int
|
29
|
+
chunk_type: str
|
30
|
+
created_at: datetime
|
31
|
+
text: str
|
32
|
+
persistent: bool = False
|
33
|
+
summary_of: Optional[List[UUID]] = None
|
34
|
+
source_chunk_id: Optional[UUID] = None
|
35
|
+
extra: Optional[Dict] = Field(default_factory=dict)
|
36
|
+
|
37
|
+
def to_chunk(self, last_hit: datetime = None, hit_count: int = 0) -> Chunk:
|
38
|
+
return Chunk(**self.model_dump(), last_hit=last_hit or datetime.now(UTC), hit_count=hit_count)
|
39
|
+
|
@@ -0,0 +1,68 @@
|
|
1
|
+
import asyncio
|
2
|
+
|
3
|
+
from google import genai
|
4
|
+
from google.genai import types
|
5
|
+
|
6
|
+
from src.lein_vector.bases.embeding_provider_abc import EmbeddingProviderABC
|
7
|
+
|
8
|
+
|
9
|
+
class EmbeddingProviderSentenceTransformer(EmbeddingProviderABC):
|
10
|
+
def __init__(self, model_name: str = "all-MiniLM-L6-v2"):
|
11
|
+
self.model = None
|
12
|
+
raise NotImplementedError
|
13
|
+
|
14
|
+
async def get_embedding(self, text: str) -> list[float]:
|
15
|
+
loop = asyncio.get_running_loop()
|
16
|
+
return await loop.run_in_executor(None, self.model.encode, text)
|
17
|
+
|
18
|
+
async def get_embeddings(self, texts: list[str]) -> list[list[float]]:
|
19
|
+
loop = asyncio.get_running_loop()
|
20
|
+
return await loop.run_in_executor(None, self.model.encode, texts)
|
21
|
+
|
22
|
+
|
23
|
+
class EmbeddingProviderGemini(EmbeddingProviderABC):
|
24
|
+
def __init__(self, api_key: str, model_name: str = "models/embedding-001"):
|
25
|
+
|
26
|
+
self.client = genai.Client(api_key=api_key)
|
27
|
+
self.model = model_name
|
28
|
+
|
29
|
+
def _genai_embed(self, text: str | list) -> types.EmbedContentResponse:
|
30
|
+
return self.client.models.embed_content(
|
31
|
+
model=self.model,
|
32
|
+
contents=text,
|
33
|
+
config=types.EmbedContentConfig(task_type="SEMANTIC_SIMILARITY"),
|
34
|
+
)
|
35
|
+
|
36
|
+
async def get_embedding(self, text: str) -> list[float]:
|
37
|
+
# В Gemini SDK обычно нет async, значит — обёртка через run_in_executor:
|
38
|
+
import asyncio
|
39
|
+
|
40
|
+
loop = asyncio.get_running_loop()
|
41
|
+
# SDK call внутри executor
|
42
|
+
embedding = await loop.run_in_executor(None, lambda: self._genai_embed(text))
|
43
|
+
# Ответ — dict с 'embedding' или 'values'
|
44
|
+
return embedding.embeddings[0].values
|
45
|
+
|
46
|
+
async def get_embeddings(self, texts: list[str]) -> list[list[float]]:
|
47
|
+
import asyncio
|
48
|
+
|
49
|
+
loop = asyncio.get_running_loop()
|
50
|
+
result = await loop.run_in_executor(None, lambda: self._genai_embed(texts))
|
51
|
+
# Вернуть список эмбеддингов
|
52
|
+
return [e.values for e in result.embeddings]
|
53
|
+
|
54
|
+
|
55
|
+
async def main():
|
56
|
+
embeding = EmbeddingProviderGemini(
|
57
|
+
api_key="AIzaSyDK_pkj25Cbb0iUujYm6N4K1k7xzeD_kss"
|
58
|
+
)
|
59
|
+
print(str(await embeding.get_embedding("test"))[:50] + "... TRIMMED]\n")
|
60
|
+
res = await embeding.get_embeddings(["test", "test2"])
|
61
|
+
print(len(res))
|
62
|
+
for e in res:
|
63
|
+
print(str(e)[:50] + "... TRIMMED]")
|
64
|
+
|
65
|
+
|
66
|
+
|
67
|
+
if __name__ == "__main__":
|
68
|
+
asyncio.run(main())
|
@@ -0,0 +1,44 @@
|
|
1
|
+
from collections import deque
|
2
|
+
from typing import List, Dict, Optional
|
3
|
+
from datetime import datetime
|
4
|
+
|
5
|
+
class ShortTermMemory:
|
6
|
+
def __init__(self, maxlen: int = 10):
|
7
|
+
self._buffer: deque = deque(maxlen=maxlen)
|
8
|
+
|
9
|
+
def add(self, role: str, text: str, ts: Optional[datetime] = None) -> None:
|
10
|
+
"""Добавить сообщение в память (роль, текст, ts - время, по умолчанию now)."""
|
11
|
+
if ts is None:
|
12
|
+
ts = datetime.now()
|
13
|
+
self._buffer.append({
|
14
|
+
"role": role,
|
15
|
+
"text": text,
|
16
|
+
"ts": ts
|
17
|
+
})
|
18
|
+
|
19
|
+
def window(self, n: Optional[int] = None) -> List[Dict]:
|
20
|
+
"""Получить последние n сообщений (по умолчанию все)."""
|
21
|
+
if n is None or n > len(self._buffer):
|
22
|
+
return list(self._buffer)
|
23
|
+
return list(self._buffer)[-n:]
|
24
|
+
|
25
|
+
def clear(self) -> None:
|
26
|
+
"""Очистить память."""
|
27
|
+
self._buffer.clear()
|
28
|
+
|
29
|
+
def load(self, history: List[Dict]) -> None:
|
30
|
+
"""Инициализировать память списком сообщений."""
|
31
|
+
self._buffer.clear()
|
32
|
+
for msg in history[-self._buffer.maxlen:]:
|
33
|
+
self._buffer.append(msg)
|
34
|
+
|
35
|
+
def to_list(self) -> List[Dict]:
|
36
|
+
"""Выгрузить всю память как список."""
|
37
|
+
return list(self._buffer)
|
38
|
+
|
39
|
+
def chunk_for_vector(self, chunk_size: int = 6) -> Optional[List[Dict]]:
|
40
|
+
"""Сформировать чанк для векторной БД — N последних сообщений по хронологии."""
|
41
|
+
if len(self._buffer) < chunk_size:
|
42
|
+
return None
|
43
|
+
# Забираем chunk_size старейших, но не очищаем (внешний код решает)
|
44
|
+
return list(self._buffer)[-chunk_size:]
|