auto embeddings

- HF embed by default
- node imports fix

Author: frdel

.vscode/launch.json

@@ -1,6 +1,7 @@
 {
   "version": "0.2.0",
   "configurations": [
+
     {
       "name": "Debug run_ui.py",
       "type": "debugpy",
@@ -10,15 +11,6 @@
       "justMyCode": false,
       "args": ["--development=true", "-Xfrozen_modules=off"]
     },
-    {
-      "name": "Debug run_cli.py",
-      "type": "debugpy",
-      "request": "launch",
-      "program": "./run_cli.py",
-      "console": "integratedTerminal",
-      "justMyCode": false,
-      "args": ["--development=true", "-Xfrozen_modules=off"]
-    },
     {
       "name": "Debug current file",
       "type": "debugpy",

agent.py

@@ -2,31 +2,21 @@
 from collections import OrderedDict
 from dataclasses import dataclass, field
 from datetime import datetime
-import time, importlib, inspect, os, json
-import token
 from typing import Any, Awaitable, Coroutine, Optional, Dict, TypedDict
 import uuid
 import models
 
-from langchain_core.prompt_values import ChatPromptValue
 from python.helpers import extract_tools, rate_limiter, files, errors, history, tokens
 from python.helpers.print_style import PrintStyle
 from langchain_core.prompts import (
     ChatPromptTemplate,
-    MessagesPlaceholder,
-    HumanMessagePromptTemplate,
-    StringPromptTemplate,
 )
-from langchain_core.prompts.image import ImagePromptTemplate
 from langchain_core.messages import HumanMessage, SystemMessage, AIMessage, BaseMessage
-from langchain_core.language_models.chat_models import BaseChatModel
-from langchain_core.language_models.llms import BaseLLM
-from langchain_core.embeddings import Embeddings
+
 import python.helpers.log as Log
 from python.helpers.dirty_json import DirtyJson
 from python.helpers.defer import DeferredTask
 from typing import Callable
-from python.helpers.history import OutputMessage
 from python.helpers.localization import Localization
 
 

docker/run/fs/exe/node_eval.js

@@ -7,8 +7,9 @@ const Module = require('module');
 // Enhance `require` to search CWD first, then globally
 function customRequire(moduleName) {
   try {
-    // Try resolving from CWD's node_modules
-    const cwdPath = path.resolve(process.cwd(), 'node_modules', moduleName);
+    // Try resolving from CWD's node_modules using Node's require.resolve
+    const cwdPath = require.resolve(moduleName, { paths: [path.join(process.cwd(), 'node_modules')] });
+    // console.log("resolved path:", cwdPath);
     return require(cwdPath);
   } catch (cwdErr) {
     try {

preload.py

@@ -1,6 +1,7 @@
 import asyncio
 from python.helpers import runtime, whisper, settings
 from python.helpers.print_style import PrintStyle
+import models
 
 PrintStyle().print("Running preload...")
 runtime.initialize()
@@ -10,12 +11,31 @@ async def preload():
     try:
         set = settings.get_default_settings()
 
+        # preload whisper model
+        async def preload_whisper():
+            try:
+                return await whisper.preload(set["stt_model_size"])
+            except Exception as e:
+                PrintStyle().error(f"Error in preload_whisper: {e}")
+
+        # preload embedding model
+        async def preload_embedding():
+            if set["embed_model_provider"] == models.ModelProvider.HUGGINGFACE.name:
+                try:
+                    emb_mod = models.get_huggingface_embedding(set["embed_model_name"])
+                    emb_txt = await emb_mod.aembed_query("test")
+                    return emb_txt
+                except Exception as e:
+                    PrintStyle().error(f"Error in preload_embedding: {e}")
+
+
         # async tasks to preload
-        tasks = [whisper.preload(set["stt_model_size"])]
+        tasks = [preload_whisper(), preload_embedding()]
 
-        return asyncio.gather(*tasks, return_exceptions=True)
+        await asyncio.gather(*tasks, return_exceptions=True)
+        PrintStyle().print("Preload completed")
     except Exception as e:
-        PrintStyle().print(f"Error in preload: {e}")
+        PrintStyle().error(f"Error in preload: {e}")
 
 
 # preload transcription model

python/helpers/files.py

@@ -285,3 +285,8 @@ def move_file(relative_path: str, new_path: str):
     new_abs_path = get_abs_path(new_path)
     os.makedirs(os.path.dirname(new_abs_path), exist_ok=True)
     os.rename(abs_path, new_abs_path)
+
+def safe_file_name(filename:str)-> str:
+    # Replace any character that's not alphanumeric, dash, underscore, or dot with underscore
+    import re
+    return re.sub(r'[^a-zA-Z0-9-._]', '_', filename)

python/helpers/memory.py

@@ -23,7 +23,7 @@
 from python.helpers import knowledge_import
 from python.helpers.log import Log, LogItem
 from enum import Enum
-from agent import Agent
+from agent import Agent, ModelConfig
 import models
 
 
@@ -36,6 +36,9 @@ def get_by_ids(self, ids: Sequence[str], /) -> List[Document]:
     async def aget_by_ids(self, ids: Sequence[str], /) -> List[Document]:
         return self.get_by_ids(ids)
 
+    def get_all_docs(self):
+        return self.docstore._dict  # type: ignore
+
 
 class Memory:
 
@@ -55,14 +58,9 @@ async def get(agent: Agent):
                 type="util",
                 heading=f"Initializing VectorDB in '/{memory_subdir}'",
             )
-            db = Memory.initialize(
+            db, created = Memory.initialize(
                 log_item,
-                models.get_model(
-                    models.ModelType.EMBEDDING,
-                    agent.config.embeddings_model.provider,
-                    agent.config.embeddings_model.name,
-                    **agent.config.embeddings_model.kwargs,
-                ),
+                agent.config.embeddings_model,
                 memory_subdir,
                 False,
             )
@@ -90,10 +88,10 @@ async def reload(agent: Agent):
     @staticmethod
     def initialize(
         log_item: LogItem | None,
-        embeddings_model: Embeddings,
+        model_config: ModelConfig,
         memory_subdir: str,
         in_memory=False,
-    ) -> MyFaiss:
+    ) -> tuple[MyFaiss, bool]:
 
         PrintStyle.standard("Initializing VectorDB...")
 
@@ -114,20 +112,26 @@ def initialize(
             os.makedirs(em_dir, exist_ok=True)
             store = LocalFileStore(em_dir)
 
+        embeddings_model = models.get_model(
+            models.ModelType.EMBEDDING,
+            model_config.provider,
+            model_config.name,
+            **model_config.kwargs,
+        )
+        embeddings_model_id = files.safe_file_name(
+            model_config.provider.name + "_" + model_config.name
+        )
+
         # here we setup the embeddings model with the chosen cache storage
         embedder = CacheBackedEmbeddings.from_bytes_store(
-            embeddings_model,
-            store,
-            namespace=getattr(
-                embeddings_model,
-                "model",
-                getattr(embeddings_model, "model_name", "default"),
-            ),
+            embeddings_model, store, namespace=embeddings_model_id
         )
 
-        # self.db = Chroma(
-        #     embedding_function=self.embedder,
-        #     persist_directory=db_dir)
+        # initial DB and docs variables
+        db: MyFaiss | None = None
+        docs: dict[str, Document] | None = None
+
+        created = False
 
         # if db folder exists and is not empty:
         if os.path.exists(db_dir) and files.exists(db_dir, "index.faiss"):
@@ -138,8 +142,27 @@ def initialize(
                 distance_strategy=DistanceStrategy.COSINE,
                 # normalize_L2=True,
                 relevance_score_fn=Memory._cosine_normalizer,
-            )
-        else:
+            )  # type: ignore
+
+            # if there is a mismatch in embeddings used, re-index the whole DB
+            emb_ok = False
+            emb_set_file = files.get_abs_path(db_dir, "embedding.json")
+            if files.exists(emb_set_file):
+                embedding_set = json.loads(files.read_file(emb_set_file))
+                if (
+                    embedding_set["model_provider"] == model_config.provider.name
+                    and embedding_set["model_name"] == model_config.name
+                ):
+                    # model matches
+                    emb_ok = True
+
+            # re-index -  create new DB and insert existing docs
+            if db and not emb_ok:
+                docs = db.get_all_docs()
+                db = None
+
+        # DB not loaded, create one
+        if not db:
             index = faiss.IndexFlatIP(len(embedder.embed_query("example")))
 
             db = MyFaiss(
@@ -151,7 +174,31 @@ def initialize(
                 # normalize_L2=True,
                 relevance_score_fn=Memory._cosine_normalizer,
             )
-        return db  # type: ignore
+
+            # insert docs if reindexing
+            if docs:
+                PrintStyle.standard("Indexing memories...")
+                if log_item:
+                    log_item.stream(progress="\nIndexing memories")
+                db.add_documents(documents=list(docs.values()), ids=list(docs.keys()))
+
+            # save DB
+            Memory._save_db_file(db, memory_subdir)
+            # save meta file
+            meta_file_path = files.get_abs_path(db_dir, "embedding.json")
+            files.write_file(
+                meta_file_path,
+                json.dumps(
+                    {
+                        "model_provider": model_config.provider.name,
+                        "model_name": model_config.name,
+                    }
+                ),
+            )
+
+            created = True
+
+        return db, created
 
     def __init__(
         self,
@@ -243,9 +290,10 @@ async def search_similarity_threshold(
     ):
         comparator = Memory._get_comparator(filter) if filter else None
 
-        #rate limiter
+        # rate limiter
         await self.agent.rate_limiter(
-            model_config=self.agent.config.embeddings_model, input=query)
+            model_config=self.agent.config.embeddings_model, input=query
+        )
 
         return await self.db.asearch(
             query,
@@ -309,25 +357,30 @@ async def insert_documents(self, docs: list[Document]):
         ids = [str(uuid.uuid4()) for _ in range(len(docs))]
         timestamp = self.get_timestamp()
 
-        
         if ids:
             for doc, id in zip(docs, ids):
                 doc.metadata["id"] = id  # add ids to documents metadata
                 doc.metadata["timestamp"] = timestamp  # add timestamp
                 if not doc.metadata.get("area", ""):
                     doc.metadata["area"] = Memory.Area.MAIN.value
-            
-            #rate limiter
+
+            # rate limiter
             docs_txt = "".join(self.format_docs_plain(docs))
             await self.agent.rate_limiter(
-                model_config=self.agent.config.embeddings_model, input=docs_txt)
+                model_config=self.agent.config.embeddings_model, input=docs_txt
+            )
 
             self.db.add_documents(documents=docs, ids=ids)
             self._save_db()  # persist
         return ids
 
     def _save_db(self):
-        self.db.save_local(folder_path=self._abs_db_dir(self.memory_subdir))
+        Memory._save_db_file(self.db, self.memory_subdir)
+
+    @staticmethod
+    def _save_db_file(db: MyFaiss, memory_subdir: str):
+        abs_dir = Memory._abs_db_dir(memory_subdir)
+        db.save_local(folder_path=abs_dir)
 
     @staticmethod
     def _get_comparator(condition: str):
@@ -382,3 +435,8 @@ def get_custom_knowledge_subdir_abs(agent: Agent) -> str:
         if dir != "default":
             return files.get_abs_path("knowledge", dir)
     raise Exception("No custom knowledge subdir set")
+
+
+def reload():
+    # clear the memory index, this will force all DBs to reload
+    Memory.index = {}