refactor: use faiss instead of mongodb for default setup

Author: maciejmajek

.gitignore

@@ -174,3 +174,5 @@ src/examples/*-demo
 artifact_database.pkl
 
 imgui.ini
+
+vectorstore_data/

config.toml

@@ -21,6 +21,10 @@ complex_model = "llama3.1:70b"
 embeddings_model = "llama3.2"
 base_url = "http://localhost:11434"
 
+[vectorstore]
+type = "faiss"
+uri = "vectorstore_data"
+
 [tracing]
 project = "rai"
 

examples/agents/spatiotemporal.py

@@ -14,7 +14,7 @@
 
 import rclpy
 from rai.agents.spatiotemporal import ROS2SpatioTemporalAgent, ROS2SpatioTemporalConfig
-from rai.utils.model_initialization import get_embeddings_model, get_llm_model
+from rai.utils.model_initialization import get_llm_model, get_vectorstore
 
 
 def create_agent():
@@ -28,7 +28,7 @@ def create_agent():
         context_compression_model=get_llm_model("simple_model"),
         time_interval=10.0,
         camera_topics=["/camera/camera/color/image_raw"],
-        embeddings=get_embeddings_model(),
+        vector_db=get_vectorstore(),
     )
     agent = ROS2SpatioTemporalAgent(config)
     return agent

src/rai_core/rai/agents/spatiotemporal/spatiotemporal_agent.py

@@ -15,12 +15,13 @@
 import json
 import logging
 import time
+import uuid
 from abc import abstractmethod
 from typing import Annotated, Any, Dict, List, Optional, cast
 
-from langchain_core.embeddings import Embeddings
 from langchain_core.language_models import BaseChatModel
 from langchain_core.messages import AIMessage, BaseMessage, HumanMessage, SystemMessage
+from langchain_core.vectorstores import VectorStore
 from pydantic import BaseModel, ConfigDict, Field
 from pymongo import MongoClient
 
@@ -60,12 +61,12 @@ class PoseStamped(BaseModel):
 
 
 class SpatioTemporalData(BaseModel):
+    id: str = Field(default_factory=lambda: str(uuid.uuid4()))
     timestamp: Annotated[float, "timestamp"]
     images: Dict[Annotated[str, "camera topic"], str] = Field(repr=False)
     tf: Optional[PoseStamped]
     temporal_context: Annotated[str, "compressed history of messages"]
     image_text_descriptions: Annotated[str, "text descriptions of images"]
-    embeddings: List[float] = Field(default_factory=list, repr=False)
 
 
 class SpatioTemporalConfig(BaseModel):
@@ -74,7 +75,7 @@ class SpatioTemporalConfig(BaseModel):
     collection_name: str
     image_to_text_model: BaseChatModel
     context_compression_model: BaseChatModel
-    embeddings: Embeddings
+    vector_db: VectorStore
     time_interval: float
 
     model_config = ConfigDict(arbitrary_types_allowed=True)
@@ -104,10 +105,9 @@ def __init__(
 
         self.db = MongoClient(self.config.db_url)[self.config.db_name]  # type: ignore
         self.collection = self.db[self.config.collection_name]  # type: ignore
-        self._initialize_embeddings_search_index()
         self.logger = logging.getLogger(__name__)
 
-    def insert_into_db(self, data: SpatioTemporalData):
+    def _insert_into_db(self, data: SpatioTemporalData):
         """
         Insert spatiotemporal data into the database.
 
@@ -121,20 +121,51 @@ def insert_into_db(self, data: SpatioTemporalData):
         )
         self.collection.insert_one(data.model_dump())  # type: ignore
 
+    def _insert_into_vectorstore(self, data: SpatioTemporalData):
+        """
+        Insert embeddings of the spatiotemporal data into the vector store.
+
+        Parameters
+        ----------
+        data : SpatioTemporalData
+            The spatiotemporal data to be inserted.
+        """
+        self.logger.info("Inserting embeddings into vector store")
+
+        print(
+            self.config.vector_db.add_texts(
+                texts=[data.temporal_context + data.image_text_descriptions],
+                metadatas=[{"id": data.id}],
+                ids=[data.id],
+            )
+        )
+
     def run(self):
         """
         Run the agent in a loop, executing tasks at specified intervals.
         """
         while True:
-            ts = time.time()
-            self.logger.info("Starting new interval")
-            self.on_interval()
-            te = time.time()
-            if te - ts > self.config.time_interval:
-                self.logger.warning(
-                    f"Time interval exceeded. Expected {self.config.time_interval:.2f}s, got {te - ts:.2f}s"
-                )
-            time.sleep(max(0, self.config.time_interval - (te - ts)))
+            try:
+                ts = time.time()
+                self.logger.info("Starting new interval")
+                self.on_interval()
+                te = time.time()
+                if te - ts > self.config.time_interval:
+                    self.logger.warning(
+                        f"Time interval exceeded. Expected {self.config.time_interval:.2f}s, got {te - ts:.2f}s"
+                    )
+                time.sleep(max(0, self.config.time_interval - (te - ts)))
+            except KeyboardInterrupt:
+                # seriously hacky
+                from langchain_community.vectorstores import FAISS
+
+                from rai.utils.model_initialization import load_config
+
+                self.logger.info("Saving vector store")
+
+                config = load_config()
+                cast(FAISS, self.config.vector_db).save_local(config.vectorstore.uri)
+                raise
 
     def on_interval(self):
         """
@@ -154,41 +185,15 @@ def on_interval(self):
         self.logger.info("Retrieving temporal context")
         temporal_context = self._get_robots_history()
 
-        embedding_text = temporal_context + str(image_text_descriptions.values())
-        self.logger.info("Embedding text")
-        embeddings = self._embed_text(embedding_text)
-
         data = SpatioTemporalData(
             timestamp=time.time(),
             images=images,
             tf=tf,
             temporal_context=temporal_context,
             image_text_descriptions=json.dumps(image_text_descriptions),
-            embeddings=embeddings,
-        )
-        self.insert_into_db(data)
-
-    def _initialize_embeddings_search_index(self):
-        self.collection.create_search_index(
-            {
-                "definition": {
-                    "mappings": {
-                        "dynamic": True,
-                        "fields": {
-                            EMBEDDINGS_FIELD_NAME: {
-                                "dimensions": 1536,
-                                "similarity": "dotProduct",
-                                "type": "knnVector",
-                            },
-                        },
-                    },
-                },
-                "name": SEARCH_INDEX_NAME,
-            }
         )
-
-    def _embed_text(self, text: str) -> List[float]:
-        return self.config.embeddings.embed_query(text)
+        self._insert_into_db(data)
+        self._insert_into_vectorstore(data)
 
     @abstractmethod
     def _get_images(

src/rai_core/rai/tools/spatiotemporal/spatiotemporal.py

@@ -13,22 +13,19 @@
 # limitations under the License.
 
 from datetime import datetime
-from typing import Any, List, Type
+from typing import Any, Dict, List, Type
 
-from langchain_core.embeddings import Embeddings
 from langchain_core.tools import BaseTool, BaseToolkit
+from langchain_core.vectorstores import VectorStore
 from pydantic import BaseModel, ConfigDict, Field
 from pymongo.collection import Collection
 from pymongo.mongo_client import MongoClient
 
 from rai.agents.spatiotemporal.spatiotemporal_agent import (
-    EMBEDDINGS_FIELD_NAME,
-    SEARCH_INDEX_NAME,
     Pose,
     SpatioTemporalData,
 )
 from rai.agents.tool_runner import MultimodalArtifact
-from rai.utils.model_initialization import get_embeddings_model
 
 
 class SpatiotemporalToolkit(BaseToolkit):
@@ -38,6 +35,7 @@ class SpatiotemporalToolkit(BaseToolkit):
     mongodb_url: str = Field(default="mongodb://localhost:27017/")
     mongodb_db_name: str = Field(default="rai")
     mongodb_collection_name: str = Field(default="spatiotemporal_collection")
+    vectorstore: VectorStore
 
     model_config = ConfigDict(arbitrary_types_allowed=True, extra="allow")
 
@@ -57,7 +55,7 @@ def get_tools(self) -> list[BaseTool]:
                 collection=self.collection,
             ),
             GetMemoriesOfObjectTool(
-                collection=self.collection,
+                collection=self.collection, vectorstore=self.vectorstore
             ),
         ]
 
@@ -205,28 +203,20 @@ class GetMemoriesOfObjectTool(BaseTool):
     description: str = "Get the past memories of the robot of a specific object"
     args_schema: Type[GetMemoriesOfObjectToolInput] = GetMemoriesOfObjectToolInput
     collection: Collection[Any]
-    embeddings: Embeddings = Field(default_factory=lambda: get_embeddings_model())
+    vectorstore: VectorStore
 
     response_model: str = "content_and_artifact"
 
     def _run(self, object_name: str, n_results: int = 5):
-        results = list(
-            self.collection.aggregate(
-                [
-                    {
-                        "$vectorSearch": {
-                            "index": SEARCH_INDEX_NAME,
-                            "path": EMBEDDINGS_FIELD_NAME,
-                            "queryVector": self.embeddings.embed_query(object_name),
-                            "numCandidates": 200,
-                            "limit": n_results,
-                        }
-                    }
-                ]
-            ),
-        )
+        documents = self.vectorstore.similarity_search(object_name, k=n_results)
+        mongodb_data: List[Dict[str, Any]] = []
+        for document in documents:
+            id = document.id
+            record = self.collection.find_one({"id": id})
+            if record is not None:
+                mongodb_data.append(record)
         images: List[str] = []
-        parsed_results = list(map(SpatioTemporalData.model_validate, results))
+        parsed_results = list(map(SpatioTemporalData.model_validate, mongodb_data))
         for result in parsed_results:
             for image in result.images.values():
                 images.append(image)

src/rai_core/rai/utils/__init__.py

@@ -11,3 +11,7 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+
+from .model_initialization import get_embeddings_model, get_llm_model, get_vectorstore
+
+__all__ = ["get_embeddings_model", "get_llm_model", "get_vectorstore"]

src/rai_core/rai/utils/model_initialization.py

@@ -73,13 +73,20 @@ class TracingConfig:
     langsmith: LangsmithConfig
 
 
+@dataclass
+class VectorStoreConfig:
+    type: str
+    uri: str
+
+
 @dataclass
 class RAIConfig:
     vendor: VendorConfig
     aws: AWSConfig
     openai: OpenAIConfig
     ollama: OllamaConfig
     tracing: TracingConfig
+    vectorstore: VectorStoreConfig
 
 
 def load_config() -> RAIConfig:
@@ -95,6 +102,7 @@ def load_config() -> RAIConfig:
             langfuse=LangfuseConfig(**config_dict["tracing"]["langfuse"]),
             langsmith=LangsmithConfig(**config_dict["tracing"]["langsmith"]),
         ),
+        vectorstore=VectorStoreConfig(**config_dict["vectorstore"]),
     )
 
 
@@ -167,6 +175,29 @@ def get_embeddings_model(vendor: str = None):
         raise ValueError(f"Unknown embeddings vendor: {vendor}")
 
 
+def get_vectorstore():
+    config = load_config()
+    logger.info(
+        f"Initializing vector store: {config.vectorstore.type} in {config.vectorstore.uri}"
+    )
+    if config.vectorstore.type == "faiss":
+        from langchain_community.vectorstores import FAISS
+
+        if os.path.exists(config.vectorstore.uri):
+            print("I EXIST")
+            return FAISS.load_local(
+                config.vectorstore.uri,
+                embeddings=get_embeddings_model(),
+                allow_dangerous_deserialization=True,
+            )
+        else:
+            index = FAISS.from_texts(["empty"], embedding=get_embeddings_model())
+            index.save_local(config.vectorstore.uri)
+            return index
+    else:
+        raise ValueError(f"Unknown vector store type: {config.vectorstore.type}")
+
+
 def get_tracing_callbacks(
     override_use_langfuse: bool = False, override_use_langsmith: bool = False
 ) -> List[BaseCallbackHandler]: