added short-term memory to the agent

Author: krik8235

.gitignore

@@ -1,8 +1,8 @@
 deploy.py
 destinations.py
 
-# knowledge/
-memory/
+entity_memory.py
+long_term_memory.py
 
 # default_agents.py
 train.py

README.md

@@ -156,14 +156,17 @@ Tasks can be delegated to a team manager, peers in the team, or completely new a
 <hr />
 
 ## Technologies Used
-**Schema, Database, Data Validation**
+**Schema, Data Validation**
    - [Pydantic](https://docs.pydantic.dev/latest/): Data validation and serialization library for Python.
    - [Pydantic_core](https://pypi.org/project/pydantic-core/): Core func packages for Pydantic.
-   - [Chroma DB](https://docs.trychroma.com/): Vector database for storing and querying usage data.
-   - [SQLite](https://www.sqlite.org/docs.html): C-language library to implements a small SQL database engine.
    - [Upstage](https://console.upstage.ai/docs/getting-started/overview): Document processer for ML tasks. (Use `Document Parser API` to extract data from documents)
    - [Docling](https://ds4sd.github.io/docling/): Document parsing
 
+**Storage**
+   - [mem0ai](https://docs.mem0.ai/quickstart#install-package): Agents' memory storage and management.
+   - [Chroma DB](https://docs.trychroma.com/): Vector database for storing and querying usage data.
+   - [SQLite](https://www.sqlite.org/docs.html): C-language library to implements a small SQL database engine.
+
 **LLM-curation**
    - [LiteLLM](https://docs.litellm.ai/docs/providers): Curation platform to access LLMs
 

pyproject.toml

@@ -45,6 +45,7 @@ dependencies = [
   "json-repair>=0.35.0",
   "wheel>=0.45.1",
   "pdfplumber>=0.11.5",
+  "mem0ai>=0.1.48",
 ]
 classifiers = [
   "Programming Language :: Python",

requirements.txt

@@ -136,11 +136,19 @@ googleapis-common-protos==1.66.0
 grpcio==1.70.0
     # via
     #   chromadb
+    #   grpcio-tools
     #   opentelemetry-exporter-otlp-proto-grpc
+    #   qdrant-client
+grpcio-tools==1.70.0
+    # via qdrant-client
 h11==0.14.0
     # via
     #   httpcore
     #   uvicorn
+h2==4.1.0
+    # via httpx
+hpack==4.1.0
+    # via h2
 httpcore==1.0.7
     # via httpx
 httptools==0.6.4
@@ -151,6 +159,7 @@ httpx==0.27.2
     #   langsmith
     #   litellm
     #   openai
+    #   qdrant-client
 huggingface-hub==0.27.0
     # via
     #   docling
@@ -159,6 +168,8 @@ huggingface-hub==0.27.0
     #   transformers
 humanfriendly==10.0
     # via coloredlogs
+hyperframe==6.1.0
+    # via h2
 idna==3.10
     # via
     #   anyio
@@ -243,6 +254,8 @@ markupsafe==3.0.2
     #   werkzeug
 mdurl==0.1.2
     # via markdown-it-py
+mem0ai==0.1.48
+    # via versionhq (pyproject.toml)
 mmh3==5.1.0
     # via chromadb
 monotonic==1.6
@@ -274,6 +287,7 @@ numpy==2.2.1
     #   onnxruntime
     #   opencv-python-headless
     #   pandas
+    #   qdrant-client
     #   safetensors
     #   scikit-image
     #   scipy
@@ -293,6 +307,7 @@ openai==1.58.1
     #   composio-openai
     #   langchain-openai
     #   litellm
+    #   mem0ai
 opencv-python-headless==4.11.0.86
     # via
     #   docling-ibm-models
@@ -377,15 +392,20 @@ pillow==10.4.0
     #   python-pptx
     #   scikit-image
     #   torchvision
+portalocker==2.10.1
+    # via qdrant-client
 posthog==3.11.0
-    # via chromadb
+    # via
+    #   chromadb
+    #   mem0ai
 propcache==0.2.1
     # via
     #   aiohttp
     #   yarl
 protobuf==5.29.3
     # via
     #   googleapis-common-protos
+    #   grpcio-tools
     #   onnxruntime
     #   opentelemetry-proto
 pyasn1==0.6.1
@@ -412,8 +432,10 @@ pydantic==2.10.6
     #   langchain-core
     #   langsmith
     #   litellm
+    #   mem0ai
     #   openai
     #   pydantic-settings
+    #   qdrant-client
 pydantic-core==2.27.2
     # via pydantic
 pydantic-settings==2.7.1
@@ -454,7 +476,9 @@ python-dotenv==1.0.1
 python-pptx==1.0.2
     # via docling
 pytz==2024.2
-    # via pandas
+    # via
+    #   mem0ai
+    #   pandas
 pyyaml==6.0.2
     # via
     #   chromadb
@@ -466,6 +490,8 @@ pyyaml==6.0.2
     #   langchain-core
     #   transformers
     #   uvicorn
+qdrant-client==1.13.2
+    # via mem0ai
 referencing==0.35.1
     # via
     #   jsonschema
@@ -528,6 +554,7 @@ sentry-sdk==2.19.2
 setuptools==75.6.0
     # via
     #   versionhq (pyproject.toml)
+    #   grpcio-tools
     #   torch
 shapely==2.0.6
     # via easyocr
@@ -546,7 +573,9 @@ sniffio==1.3.1
 soupsieve==2.6
     # via beautifulsoup4
 sqlalchemy==2.0.36
-    # via langchain
+    # via
+    #   langchain
+    #   mem0ai
 starlette==0.41.3
     # via fastapi
 sympy==1.13.3
@@ -627,6 +656,7 @@ tzdata==2025.1
 urllib3==2.3.0
     # via
     #   kubernetes
+    #   qdrant-client
     #   requests
     #   sentry-sdk
     #   types-requests

src/versionhq/agent/model.py

@@ -12,6 +12,8 @@
 from versionhq.llm.model import LLM, DEFAULT_CONTEXT_WINDOW_SIZE, DEFAULT_MODEL_NAME
 from versionhq.tool.model import Tool, ToolSet
 from versionhq.knowledge.model import BaseKnowledgeSource, Knowledge
+from versionhq.memory.contextual_memory import ContextualMemory
+from versionhq.memory.model import ShortTermMemory, UserMemory
 from versionhq._utils.logger import Logger
 from versionhq.agent.rpm_controller import RPMController
 from versionhq._utils.usage_metrics import UsageMetrics
@@ -95,10 +97,20 @@ class Agent(BaseModel):
     backstory: Optional[str] = Field(default=None, description="developer prompt to the llm")
     skillsets: Optional[List[str]] = Field(default_factory=list)
     tools: Optional[List[Tool | ToolSet | Type[Tool]]] = Field(default_factory=list)
+
+    # knowledge
     knowledge_sources: Optional[List[BaseKnowledgeSource]] = Field(default=None)
     _knowledge: Optional[Knowledge] = PrivateAttr(default=None)
-    embedder_config: Optional[Dict[str, Any]] = Field(default=None, description="embedder configuration for the agent's knowledge")
 
+    # memory
+    use_memory: bool = Field(default=False, description="whether to store/use memory when executing the task")
+    memory_config: Optional[Dict[str, Any]] = Field(default=None, description="configuration for the memory")
+    short_term_memory: Optional[InstanceOf[ShortTermMemory]] = Field(default=None)
+    user_memory: Optional[InstanceOf[UserMemory]] = Field(default=None)
+    # _short_term_memory: Optional[InstanceOf[ShortTermMemory]] = PrivateAttr()
+    # _user_memory: Optional[InstanceOf[UserMemory]] = PrivateAttr()
+
+    embedder_config: Optional[Dict[str, Any]] = Field(default=None, description="embedder configuration for the agent's knowledge")
 
     # prompting
     use_developer_prompt: Optional[bool] = Field(default=True, description="Use developer prompt when calling the llm")
@@ -347,14 +359,30 @@ def set_up_knowledge(self) -> Self:
         return self
 
 
+    @model_validator(mode="after")
+    def set_up_memory(self) -> Self:
+        """
+        Set up memories: stm, um
+        """
+
+        if self.use_memory == True:
+            self.short_term_memory = self.short_term_memory if self.short_term_memory else ShortTermMemory(agent=self, embedder_config=self.embedder_config)
+
+            if hasattr(self, "memory_config") and self.memory_config is not None:
+                self.user_memory = self.user_memory if self.user_memory else UserMemory(agent=self)
+            else:
+                self.user_memory = None
+
+        return self
+
+
     def _train(self) -> Self:
         """
         Fine-tuned the base model using OpenAI train framework.
         """
         if not isinstance(self.llm, LLM):
             pass
 
-
     def invoke(
         self,
         prompts: str,
@@ -440,6 +468,14 @@ def execute_task(self, task, context: Optional[str] = None, task_tools: Optional
                 if agent_knowledge_context:
                     task_prompt += agent_knowledge_context
 
+
+        if self.use_memory == True:
+            contextual_memory = ContextualMemory(memory_config=self.memory_config, stm=self.short_term_memory, um=self.user_memory)
+            memory = contextual_memory.build_context_for_task(task, context)
+            if memory.strip() != "":
+                task_prompt += memory.strip()
+
+
         # if self.team and self.team._train:
         #     task_prompt = self._training_handler(task_prompt=task_prompt)
         # else:

src/versionhq/memory/__init__.py

@@ -0,0 +1,100 @@
+from typing import Any, Dict, Optional, List
+
+from versionhq.memory.model import ShortTermMemory, UserMemory
+
+
+class ContextualMemory:
+    """
+    A class to construct context from memories (ShortTermMemory, UserMemory).
+    The context will be added to the prompt when the agent executes the task.
+    """
+
+    def __init__(
+        self,
+        memory_config: Optional[Dict[str, Any]],
+        stm: ShortTermMemory,
+        um: UserMemory,
+        # ltm: LongTermMemory,
+        # em: EntityMemory,
+    ):
+        self.memory_provider = memory_config.get("provider") if memory_config is not None else None
+        self.stm = stm
+        self.um = um
+
+
+    def build_context_for_task(self, task, context: List[Any] | str) -> str:
+        """
+        Automatically builds a minimal, highly relevant set of contextual information for a given task.
+        """
+
+        query = f"{task.description} {context}".strip()
+
+        if query == "":
+            return ""
+
+        context = []
+        context.append(self._fetch_stm_context(query))
+        if self.memory_provider == "mem0":
+            context.append(self._fetch_user_context(query))
+        return "\n".join(filter(None, context))
+
+
+    def _fetch_stm_context(self, query) -> str:
+        """
+        Fetches recent relevant insights from STM related to the task's description and expected_output, formatted as bullet points.
+        """
+        stm_results = self.stm.search(query)
+        formatted_results = "\n".join(
+            [
+                f"- {result['memory'] if self.memory_provider == 'mem0' else result['context']}"
+                for result in stm_results
+            ]
+        )
+        return f"Recent Insights:\n{formatted_results}" if stm_results else ""
+
+
+    def _fetch_user_context(self, query: str) -> str:
+        """
+        Fetches and formats relevant user information from User Memory.
+        """
+
+        user_memories = self.um.search(query)
+        if not user_memories:
+            return ""
+
+        formatted_memories = "\n".join(f"- {result['memory']}" for result in user_memories)
+        return f"User memories/preferences:\n{formatted_memories}"
+
+
+    # def _fetch_ltm_context(self, task) -> Optional[str]:
+    #     """
+    #     Fetches historical data or insights from LTM that are relevant to the task's description and expected_output,
+    #     formatted as bullet points.
+    #     """
+    #     ltm_results = self.ltm.search(task, latest_n=2)
+    #     if not ltm_results:
+    #         return None
+
+    #     formatted_results = [
+    #         suggestion
+    #         for result in ltm_results
+    #         for suggestion in result["metadata"]["suggestions"]  # type: ignore # Invalid index type "str" for "str"; expected type "SupportsIndex | slice"
+    #     ]
+    #     formatted_results = list(dict.fromkeys(formatted_results))
+    #     formatted_results = "\n".join([f"- {result}" for result in formatted_results])  # type: ignore # Incompatible types in assignment (expression has type "str", variable has type "list[str]")
+
+    #     return f"Historical Data:\n{formatted_results}" if ltm_results else ""
+
+    # def _fetch_entity_context(self, query) -> str:
+    #     """
+    #     Fetches relevant entity information from Entity Memory related to the task's description and expected_output,
+    #     formatted as bullet points.
+    #     """
+    #     em_results = self.em.search(query)
+    #     formatted_results = "\n".join(
+    #         [
+    #             f"- {result['memory'] if self.memory_provider == 'mem0' else result['context']}"
+    #             for result in em_results
+    #         ]  # type: ignore #  Invalid index type "str" for "str"; expected type "SupportsIndex | slice"
+    #     )
+    #     return f"Entities:\n{formatted_results}" if em_results else ""