chore: cleanup

docs: update
type: ignore

docs: add diagram

Author: maciejmajek

docs/agents/spatiotemporal.md

@@ -1,5 +1,7 @@
 # SpatioTemporalAgent
 
+![SpatioTemporalAgent](../imgs/spatiotemporal.png)
+
 ## Overview
 
 The `SpatioTemporalAgent` is an intelligent agent designed to capture, process, and store spatiotemporal data - information that combines both spatial (location, pose) and temporal (time-based) aspects. It's particularly useful in robotics and autonomous systems where understanding both the spatial context and its evolution over time is crucial.
@@ -34,7 +36,7 @@ The agent is configured through `SpatioTemporalConfig`:
 - `collection_name`: Collection for data storage
 - `image_to_text_model`: AI model for image description
 - `context_compression_model`: AI model for context compression
-- `embeddings`: Embeddings model
+- `vector_db`: Vector store
 - `time_interval`: Data collection frequency
 
 ## ROS2 Implementation
@@ -62,20 +64,21 @@ docker run -d --name rai-mongo -p 27017:27017 mongo
 2. **Agent Configuration**
 
 ```python
+import rclpy
 from rai.agents.spatiotemporal import ROS2SpatioTemporalAgent, ROS2SpatioTemporalConfig
-from rai.utils.model_initialization import get_llm_model, get_embeddings_model
+from rai.utils.model_initialization import get_llm_model, get_vectorstore
 
 config = ROS2SpatioTemporalConfig(
     robot_frame="base_link",
     world_frame="world",
     db_url="mongodb://localhost:27017/",
     db_name="rai",
-    collection_name="spatiotemporal_data",
+    collection_name="spatiotemporal_collection",
     image_to_text_model=get_llm_model("simple_model"),
     context_compression_model=get_llm_model("simple_model"),
     time_interval=10.0,
-    camera_topics=["/camera/color/image_raw"],
-    embeddings=get_embeddings_model(),
+    camera_topics=["/camera/camera/color/image_raw"],
+    vector_db=get_vectorstore(),
 )
 
 agent = ROS2SpatioTemporalAgent(config)

docs/imgs/spatiotemporal.png

@@ -37,7 +37,7 @@ class ROS2SpatioTemporalConfig(SpatioTemporalConfig):
 
 
 class ROS2SpatioTemporalAgent(SpatioTemporalAgent):
-    def __init__(self, config: ROS2SpatioTemporalConfig, *args, **kwargs):
+    def __init__(self, config: ROS2SpatioTemporalConfig, *args, **kwargs):  # type: ignore
         super().__init__(config, *args, **kwargs)
         self.config = config
         self.connector = ROS2ARIConnector()
@@ -48,7 +48,7 @@ def _get_images(self) -> Dict[Annotated[str, "camera topic"], str]:
         images: Dict[Annotated[str, "camera topic"], str] = {}
         for camera_topic in self.config.camera_topics:
             try:
-                _, artifact = self.get_image_tool._run(topic=camera_topic)
+                _, artifact = self.get_image_tool._run(topic=camera_topic)  # type: ignore
                 image = artifact["images"][0]
                 images[camera_topic] = image
             except Exception as e:
@@ -67,21 +67,21 @@ def _get_tf(self) -> Optional[PoseStamped]:
             return None
         ps = PoseStamped(
             header=Header(
-                stamp=tf_stamped.header.stamp.sec
-                + tf_stamped.header.stamp.nanosec * 1e-9,
-                frame_id=tf_stamped.header.frame_id,
+                stamp=tf_stamped.header.stamp.sec  # type: ignore
+                + tf_stamped.header.stamp.nanosec * 1e-9,  # type: ignore
+                frame_id=tf_stamped.header.frame_id,  # type: ignore
             ),
             pose=Pose(
                 position=Point(
-                    x=tf_stamped.transform.translation.x,
-                    y=tf_stamped.transform.translation.y,
-                    z=tf_stamped.transform.translation.z,
+                    x=tf_stamped.transform.translation.x,  # type: ignore
+                    y=tf_stamped.transform.translation.y,  # type: ignore
+                    z=tf_stamped.transform.translation.z,  # type: ignore
                 ),
                 orientation=Quaternion(
-                    x=tf_stamped.transform.rotation.x,
-                    y=tf_stamped.transform.rotation.y,
-                    z=tf_stamped.transform.rotation.z,
-                    w=tf_stamped.transform.rotation.w,
+                    x=tf_stamped.transform.rotation.x,  # type: ignore
+                    y=tf_stamped.transform.rotation.y,  # type: ignore
+                    z=tf_stamped.transform.rotation.z,  # type: ignore
+                    w=tf_stamped.transform.rotation.w,  # type: ignore
                 ),
             ),
         )

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

@@ -28,9 +28,6 @@
 from rai.agents.base import BaseAgent
 from rai.messages.multimodal import HumanMultimodalMessage
 
-EMBEDDINGS_FIELD_NAME = "embeddings"
-SEARCH_INDEX_NAME = "embeddings_search_index"
-
 
 class Header(BaseModel):
     stamp: Annotated[float, "timestamp"]
@@ -100,7 +97,7 @@ def __init__(
         **kwargs : dict
             Additional keyword arguments.
         """
-        super().__init__(*args, **kwargs)
+        super().__init__(*args, **kwargs)  # type: ignore
         self.config = config
 
         self.db = MongoClient(self.config.db_url)[self.config.db_name]  # type: ignore
@@ -133,7 +130,7 @@ def _insert_into_vectorstore(self, data: SpatioTemporalData):
         self.logger.info("Inserting embeddings into vector store")
 
         print(
-            self.config.vector_db.add_texts(
+            self.config.vector_db.add_texts(  # type: ignore
                 texts=[data.temporal_context + data.image_text_descriptions],
                 metadatas=[{"id": data.id}],
                 ids=[data.id],
@@ -252,7 +249,7 @@ def _get_image_text_descriptions(
                     [text_description_prompt, human_message]
                 ),
             )
-            if not isinstance(ai_msg.content, str):
+            if not isinstance(ai_msg.content, str):  # type: ignore
                 raise ValueError("AI message content is not a string")
             text_descriptions[source] = ai_msg.content
 
@@ -289,7 +286,8 @@ def _compress_context(self, history: List[BaseMessage]) -> str:
         )
 
         robots_history: List[Dict[str, str]] = [
-            {"role": msg.type, "content": msg.content} for msg in history
+            {"role": msg.type, "content": msg.content}  # type: ignore
+            for msg in history
         ]
         if len(robots_history) == 0:
             return ""
@@ -301,6 +299,6 @@ def _compress_context(self, history: List[BaseMessage]) -> str:
                 [system_prompt, human_message]
             ),
         )
-        if not isinstance(ai_msg.content, str):
+        if not isinstance(ai_msg.content, str):  # type: ignore
             raise ValueError("AI message content is not a string")
         return ai_msg.content

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

@@ -15,7 +15,7 @@
 import logging
 import os
 from dataclasses import dataclass
-from typing import List, Literal, Optional, cast
+from typing import Any, List, Literal, Optional, cast
 
 import coloredlogs
 import tomli
@@ -109,7 +109,7 @@ def load_config() -> RAIConfig:
 def get_llm_model(
     model_type: Literal["simple_model", "complex_model"],
     vendor: Optional[str] = None,
-    **kwargs,
+    **kwargs: Any,
 ):
     config = load_config()
     if vendor is None:
@@ -184,7 +184,6 @@ def get_vectorstore():
         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(),