auto embeddings
- HF embed by default - node imports fix
frdel committed
May 6, 2025 at 16:21 UTC
a2c03c06d95deec91455c4f2c3c9fa1a326f2d38
7 files changed
+138
-63
.vscode/launch.json
+1
-9
@@ -1,6 +1,7 @@
1
{
2
"version": "0.2.0",
3
"configurations": [
4
+
5
{
6
"name": "Debug run_ui.py",
7
"type": "debugpy",
@@ -10,15 +11,6 @@
11
"justMyCode": false,
12
"args": ["--development=true", "-Xfrozen_modules=off"]
13
},
13
- {
14
- "name": "Debug run_cli.py",
15
- "type": "debugpy",
16
- "request": "launch",
17
- "program": "./run_cli.py",
18
- "console": "integratedTerminal",
19
- "justMyCode": false,
20
- "args": ["--development=true", "-Xfrozen_modules=off"]
21
- },
14
{
15
"name": "Debug current file",
16
"type": "debugpy",
agent.py
+1
-11
@@ -2,31 +2,21 @@ import asyncio
2
from collections import OrderedDict
3
from dataclasses import dataclass, field
4
from datetime import datetime
5
-import time, importlib, inspect, os, json
6
-import token
5
from typing import Any, Awaitable, Coroutine, Optional, Dict, TypedDict
6
import uuid
7
import models
8
11
-from langchain_core.prompt_values import ChatPromptValue
9
from python.helpers import extract_tools, rate_limiter, files, errors, history, tokens
10
from python.helpers.print_style import PrintStyle
11
from langchain_core.prompts import (
12
ChatPromptTemplate,
16
- MessagesPlaceholder,
17
- HumanMessagePromptTemplate,
18
- StringPromptTemplate,
13
)
20
-from langchain_core.prompts.image import ImagePromptTemplate
14
from langchain_core.messages import HumanMessage, SystemMessage, AIMessage, BaseMessage
22
-from langchain_core.language_models.chat_models import BaseChatModel
23
-from langchain_core.language_models.llms import BaseLLM
24
-from langchain_core.embeddings import Embeddings
15
+
16
import python.helpers.log as Log
17
from python.helpers.dirty_json import DirtyJson
18
from python.helpers.defer import DeferredTask
19
from typing import Callable
29
-from python.helpers.history import OutputMessage
20
from python.helpers.localization import Localization
21
22
docker/run/fs/exe/node_eval.js
+3
-2
@@ -7,8 +7,9 @@ const Module = require('module');
7
// Enhance `require` to search CWD first, then globally
8
function customRequire(moduleName) {
9
try {
10
- // Try resolving from CWD's node_modules
11
- const cwdPath = path.resolve(process.cwd(), 'node_modules', moduleName);
10
+ // Try resolving from CWD's node_modules using Node's require.resolve
11
+ const cwdPath = require.resolve(moduleName, { paths: [path.join(process.cwd(), 'node_modules')] });
12
+ // console.log("resolved path:", cwdPath);
13
return require(cwdPath);
14
} catch (cwdErr) {
15
try {
preload.py
+23
-3
@@ -1,6 +1,7 @@
1
import asyncio
2
from python.helpers import runtime, whisper, settings
3
from python.helpers.print_style import PrintStyle
4
+import models
5
6
PrintStyle().print("Running preload...")
7
runtime.initialize()
@@ -10,12 +11,31 @@ async def preload():
11
try:
12
set = settings.get_default_settings()
13
14
+ # preload whisper model
15
+ async def preload_whisper():
16
+ try:
17
+ return await whisper.preload(set["stt_model_size"])
18
+ except Exception as e:
19
+ PrintStyle().error(f"Error in preload_whisper: {e}")
20
+
21
+ # preload embedding model
22
+ async def preload_embedding():
23
+ if set["embed_model_provider"] == models.ModelProvider.HUGGINGFACE.name:
24
+ try:
25
+ emb_mod = models.get_huggingface_embedding(set["embed_model_name"])
26
+ emb_txt = await emb_mod.aembed_query("test")
27
+ return emb_txt
28
+ except Exception as e:
29
+ PrintStyle().error(f"Error in preload_embedding: {e}")
30
+
31
+
32
# async tasks to preload
14
- tasks = [whisper.preload(set["stt_model_size"])]
33
+ tasks = [preload_whisper(), preload_embedding()]
34
16
- return asyncio.gather(*tasks, return_exceptions=True)
35
+ await asyncio.gather(*tasks, return_exceptions=True)
36
+ PrintStyle().print("Preload completed")
37
except Exception as e:
18
- PrintStyle().print(f"Error in preload: {e}")
38
+ PrintStyle().error(f"Error in preload: {e}")
39
40
41
# preload transcription model
python/helpers/files.py
+5
@@ -285,3 +285,8 @@ def move_file(relative_path: str, new_path: str):
285
new_abs_path = get_abs_path(new_path)
286
os.makedirs(os.path.dirname(new_abs_path), exist_ok=True)
287
os.rename(abs_path, new_abs_path)
288
+
289
+def safe_file_name(filename:str)-> str:
290
+ # Replace any character that's not alphanumeric, dash, underscore, or dot with underscore
291
+ import re
292
+ return re.sub(r'[^a-zA-Z0-9-._]', '_', filename)
\ No newline at end of file
python/helpers/memory.py
+88
-30
@@ -23,7 +23,7 @@ import uuid
23
from python.helpers import knowledge_import
24
from python.helpers.log import Log, LogItem
25
from enum import Enum
26
-from agent import Agent
26
+from agent import Agent, ModelConfig
27
import models
28
29
@@ -36,6 +36,9 @@ class MyFaiss(FAISS):
36
async def aget_by_ids(self, ids: Sequence[str], /) -> List[Document]:
37
return self.get_by_ids(ids)
38
39
+ def get_all_docs(self):
40
+ return self.docstore._dict # type: ignore
41
+
42
43
class Memory:
44
@@ -55,14 +58,9 @@ class Memory:
58
type="util",
59
heading=f"Initializing VectorDB in '/{memory_subdir}'",
60
)
58
- db = Memory.initialize(
61
+ db, created = Memory.initialize(
62
log_item,
60
- models.get_model(
61
- models.ModelType.EMBEDDING,
62
- agent.config.embeddings_model.provider,
63
- agent.config.embeddings_model.name,
64
- **agent.config.embeddings_model.kwargs,
65
- ),
63
+ agent.config.embeddings_model,
64
memory_subdir,
65
False,
66
)
@@ -90,10 +88,10 @@ class Memory:
88
@staticmethod
89
def initialize(
90
log_item: LogItem | None,
93
- embeddings_model: Embeddings,
91
+ model_config: ModelConfig,
92
memory_subdir: str,
93
in_memory=False,
96
- ) -> MyFaiss:
94
+ ) -> tuple[MyFaiss, bool]:
95
96
PrintStyle.standard("Initializing VectorDB...")
97
@@ -114,20 +112,26 @@ class Memory:
112
os.makedirs(em_dir, exist_ok=True)
113
store = LocalFileStore(em_dir)
114
115
+ embeddings_model = models.get_model(
116
+ models.ModelType.EMBEDDING,
117
+ model_config.provider,
118
+ model_config.name,
119
+ **model_config.kwargs,
120
+ )
121
+ embeddings_model_id = files.safe_file_name(
122
+ model_config.provider.name + "_" + model_config.name
123
+ )
124
+
125
# here we setup the embeddings model with the chosen cache storage
126
embedder = CacheBackedEmbeddings.from_bytes_store(
119
- embeddings_model,
120
- store,
121
- namespace=getattr(
122
- embeddings_model,
123
- "model",
124
- getattr(embeddings_model, "model_name", "default"),
125
- ),
127
+ embeddings_model, store, namespace=embeddings_model_id
128
)
129
128
- # self.db = Chroma(
129
- # embedding_function=self.embedder,
130
- # persist_directory=db_dir)
130
+ # initial DB and docs variables
131
+ db: MyFaiss | None = None
132
+ docs: dict[str, Document] | None = None
133
+
134
+ created = False
135
136
# if db folder exists and is not empty:
137
if os.path.exists(db_dir) and files.exists(db_dir, "index.faiss"):
@@ -138,8 +142,27 @@ class Memory:
142
distance_strategy=DistanceStrategy.COSINE,
143
# normalize_L2=True,
144
relevance_score_fn=Memory._cosine_normalizer,
141
- )
142
- else:
145
+ ) # type: ignore
146
+
147
+ # if there is a mismatch in embeddings used, re-index the whole DB
148
+ emb_ok = False
149
+ emb_set_file = files.get_abs_path(db_dir, "embedding.json")
150
+ if files.exists(emb_set_file):
151
+ embedding_set = json.loads(files.read_file(emb_set_file))
152
+ if (
153
+ embedding_set["model_provider"] == model_config.provider.name
154
+ and embedding_set["model_name"] == model_config.name
155
+ ):
156
+ # model matches
157
+ emb_ok = True
158
+
159
+ # re-index - create new DB and insert existing docs
160
+ if db and not emb_ok:
161
+ docs = db.get_all_docs()
162
+ db = None
163
+
164
+ # DB not loaded, create one
165
+ if not db:
166
index = faiss.IndexFlatIP(len(embedder.embed_query("example")))
167
168
db = MyFaiss(
@@ -151,7 +174,31 @@ class Memory:
174
# normalize_L2=True,
175
relevance_score_fn=Memory._cosine_normalizer,
176
)
154
- return db # type: ignore
177
+
178
+ # insert docs if reindexing
179
+ if docs:
180
+ PrintStyle.standard("Indexing memories...")
181
+ if log_item:
182
+ log_item.stream(progress="\nIndexing memories")
183
+ db.add_documents(documents=list(docs.values()), ids=list(docs.keys()))
184
+
185
+ # save DB
186
+ Memory._save_db_file(db, memory_subdir)
187
+ # save meta file
188
+ meta_file_path = files.get_abs_path(db_dir, "embedding.json")
189
+ files.write_file(
190
+ meta_file_path,
191
+ json.dumps(
192
+ {
193
+ "model_provider": model_config.provider.name,
194
+ "model_name": model_config.name,
195
+ }
196
+ ),
197
+ )
198
+
199
+ created = True
200
+
201
+ return db, created
202
203
def __init__(
204
self,
@@ -243,9 +290,10 @@ class Memory:
290
):
291
comparator = Memory._get_comparator(filter) if filter else None
292
246
- #rate limiter
293
+ # rate limiter
294
await self.agent.rate_limiter(
248
- model_config=self.agent.config.embeddings_model, input=query)
295
+ model_config=self.agent.config.embeddings_model, input=query
296
+ )
297
298
return await self.db.asearch(
299
query,
@@ -309,25 +357,30 @@ class Memory:
357
ids = [str(uuid.uuid4()) for _ in range(len(docs))]
358
timestamp = self.get_timestamp()
359
312
-
360
if ids:
361
for doc, id in zip(docs, ids):
362
doc.metadata["id"] = id # add ids to documents metadata
363
doc.metadata["timestamp"] = timestamp # add timestamp
364
if not doc.metadata.get("area", ""):
365
doc.metadata["area"] = Memory.Area.MAIN.value
319
-
320
- #rate limiter
366
+
367
+ # rate limiter
368
docs_txt = "".join(self.format_docs_plain(docs))
369
await self.agent.rate_limiter(
323
- model_config=self.agent.config.embeddings_model, input=docs_txt)
370
+ model_config=self.agent.config.embeddings_model, input=docs_txt
371
+ )
372
373
self.db.add_documents(documents=docs, ids=ids)
374
self._save_db() # persist
375
return ids
376
377
def _save_db(self):
330
- self.db.save_local(folder_path=self._abs_db_dir(self.memory_subdir))
378
+ Memory._save_db_file(self.db, self.memory_subdir)
379
+
380
+ @staticmethod
381
+ def _save_db_file(db: MyFaiss, memory_subdir: str):
382
+ abs_dir = Memory._abs_db_dir(memory_subdir)
383
+ db.save_local(folder_path=abs_dir)
384
385
@staticmethod
386
def _get_comparator(condition: str):
@@ -382,3 +435,8 @@ def get_custom_knowledge_subdir_abs(agent: Agent) -> str:
435
if dir != "default":
436
return files.get_abs_path("knowledge", dir)
437
raise Exception("No custom knowledge subdir set")
438
+
439
+
440
+def reload():
441
+ # clear the memory index, this will force all DBs to reload
442
+ Memory.index = {}
python/helpers/settings.py
+17
-8
@@ -150,7 +150,6 @@ def convert_out(settings: Settings) -> SettingsOutput:
150
}
151
)
152
153
-
153
chat_model_fields.append(
154
{
155
"id": "chat_model_vision",
@@ -730,9 +729,10 @@ def get_settings() -> Settings:
729
730
def set_settings(settings: Settings):
731
global _settings
732
+ previous = _settings
733
_settings = normalize_settings(settings)
734
_write_settings_file(_settings)
735
- _apply_settings()
735
+ _apply_settings(previous)
736
737
738
def normalize_settings(settings: Settings) -> Settings:
@@ -795,7 +795,7 @@ def get_default_settings() -> Settings:
795
return Settings(
796
chat_model_provider=ModelProvider.OPENAI.name,
797
chat_model_name="gpt-4o",
798
- chat_model_kwargs={ "temperature": "0" },
798
+ chat_model_kwargs={"temperature": "0"},
799
chat_model_ctx_length=120000,
800
chat_model_ctx_history=0.7,
801
chat_model_vision=False,
@@ -806,19 +806,19 @@ def get_default_settings() -> Settings:
806
util_model_name="gpt-4o-mini",
807
util_model_ctx_length=120000,
808
util_model_ctx_input=0.7,
809
- util_model_kwargs={ "temperature": "0" },
809
+ util_model_kwargs={"temperature": "0"},
810
util_model_rl_requests=60,
811
util_model_rl_input=0,
812
util_model_rl_output=0,
813
- embed_model_provider=ModelProvider.OPENAI.name,
814
- embed_model_name="text-embedding-3-small",
813
+ embed_model_provider=ModelProvider.HUGGINGFACE.name,
814
+ embed_model_name="sentence-transformers/all-MiniLM-L6-v2",
815
embed_model_kwargs={},
816
embed_model_rl_requests=0,
817
embed_model_rl_input=0,
818
browser_model_provider=ModelProvider.OPENAI.name,
819
browser_model_name="gpt-4o",
820
browser_model_vision=False,
821
- browser_model_kwargs={ "temperature": "0" },
821
+ browser_model_kwargs={"temperature": "0"},
822
api_keys={},
823
auth_login="",
824
auth_password="",
@@ -839,7 +839,7 @@ def get_default_settings() -> Settings:
839
)
840
841
842
-def _apply_settings():
842
+def _apply_settings(previous: Settings | None):
843
global _settings
844
if _settings:
845
from agent import AgentContext
@@ -858,6 +858,15 @@ def _apply_settings():
858
whisper.preload, _settings["stt_model_size"]
859
) # TODO overkill, replace with background task
860
861
+ # force memory reload on embedding model change
862
+ if previous and (
863
+ _settings["embed_model_name"] != previous["embed_model_name"]
864
+ or _settings["embed_model_provider"] != previous["embed_model_provider"]
865
+ or _settings["embed_model_kwargs"] != previous["embed_model_kwargs"]
866
+ ):
867
+ from python.helpers.memory import reload as memory_reload
868
+ memory_reload()
869
+
870
871
def _env_to_dict(data: str):
872
env_dict = {}