You are viewing a plain text version of this content. The canonical link for it is here.
Posted to issues@flink.apache.org by "wenjin272 (via GitHub)" <gi...@apache.org> on 2025/11/20 08:11:19 UTC
[PR] Long term memory [flink-agents]
wenjin272 opened a new pull request, #332:
URL: https://github.com/apache/flink-agents/pull/332
<!--
* Thank you very much for contributing to Flink Agents.
* Please add the relevant components in the PR title. E.g., [api], [runtime], [java], [python], [hotfix], etc.
-->
<!-- Please link the PR to the relevant issue(s). Hotfix doesn't need this. -->
Linked issue: #331
### Purpose of change
<!-- What is the purpose of this change? -->
Introduce the long-term memory interface in python, and provide an implementation based on chroma.
This is the first pr of three to introduce long-term memory in python:
1. interface and one implementation
2. support using long-term memory in action
3. async interface and execution
### Tests
Unit test
### API
Yes, add long-term memory related api.
### Documentation
<!-- Do not remove this section. Check the proper box only. -->
- [ ] `doc-needed` <!-- Your PR changes impact docs -->
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@flink.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] [api][runtime] Introduce long-term memory in python [flink-agents]
Posted by "alnzng (via GitHub)" <gi...@apache.org>.
alnzng commented on PR #332:
URL: https://github.com/apache/flink-agents/pull/332#issuecomment-3604484837
> Hi, @alnzng. There's a design issue related to the vector store that I'd appreciate your help reviewing.
>
> As describe in the design doc #339, long-term memory of flink-agents is also based on vector store. Currently, I provide an implementation based on chroma. In this implementation, I directly use chroma client rather than flink-agents `BaseVectorStore`, because there are some long-term memory needed interface not provided in `BaseVectorStore`.
>
> @xintongsong believes that we can directly build upon the Flink-Agents `BaseVectorStore`. Thus, we can support using any already supported vector store as the backend for long-term memory.
>
> I think this make sense, but it requires add some interfaces to `BaseVectorStore`, which maybe look like:
>
> ```
> def get_or_create_collection(self, name: str, metadata: Dict[str, Any]) -> Dict[str, Any]:
> """Get a collection, create if it doesn't already exist."""
>
> def get_collection(self, name: str) -> None:
> """Get an existing collection."""
>
> def update_collection(self, name: str, metadata: Dict[str, Any]) -> None:
> """Update an existing collection."""
>
> def delete_collection(self, name: str) -> bool:
> """Delete a collection."""
>
> def add(self, document: Document, collection_name: str | None = None) -> None:
> """Add a document to the collection."""
>
> def update(self, document: Document, collection_name: str | None = None) -> None:
> """Update a document, can only update metadata."""
>
> def delete(self, offset: int | None, limit: int | None, ids: List[int] | None = None, **kwargs: Any) -> bool:
> """Delete documents from collection."""
>
> def get(self, offset: int | None, limit: int | None, collection_name: str | None = None, **kwargs: Any) -> List[Document]:
> """Get documents from collection."""
> ```
>
> These interface may not be achievable for each vector store, I will conduct research and refinement afterward. WDTY?
Thanks for raising this design question @wenjin272! I agree with @xintongsong that building long-term memory on top of `BaseVectorStore` would be beneficial for consistency and reusability.
The current `BaseVectorStore` is intentionally minimal, focusing on read-only semantic search. This design works well for RAG retrieval use cases but indeed lacks the CRUD operations needed for long-term memory management. Adding document-level CRUD operations to `BaseVectorStore` makes sense, as these operations are generally supported across all major vector stores.
A few thoughts on the proposed methods:
- add(), delete(), get() - Widely supported, should be included
- update() - I'm not certain about common use cases for updating existing long-term memory entries, but it shouldn't be
problematic to include.
However, I have concerns about including collection management in the base interface. This concept is vector
store-specific and may cause integration issues. "Collection" is most common term (e.g. Chroma, Milvus, Qdrant, etc), but Pinecone and Weaviate use different terms(e.g. Class, Namespace, etc).
My suggestion: Keep collection/namespace management implementation-specific rather than in the base interface. Each vector store integration can expose collection management using its native terminology and patterns.
This approach aligns with how `LlamaIndex` and `LangChain` handle vector stores. I'd recommend you reviewing these implementations, they provide good precedents for balancing abstraction with flexibility.
- LlamaIndex BasePydanticVectorStore:
https://github.com/run-llama/llama_index/blob/main/llama-index-core/llama_index/core/vector_stores/types.py#L334
- LangChain VectorStore: https://github.com/langchain-ai/langchain/blob/master/libs/core/langchain_core/vectorstores/base.py#L43
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@flink.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] [api][runtime] Introduce long-term memory in python [flink-agents]
Posted by "wenjin272 (via GitHub)" <gi...@apache.org>.
wenjin272 commented on code in PR #332:
URL: https://github.com/apache/flink-agents/pull/332#discussion_r2602172199
##########
python/flink_agents/runtime/memory/tests/start_chroma_server.sh:
##########
@@ -0,0 +1,23 @@
+#!/usr/bin/env bash
+
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# 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.
+#
+#
+
+path=$1
+chroma run --path $path
Review Comment:
IIUC, the `import chromadb` will start the chroma in In-Memory mode. For long-term memory, we recommend Server-Client mode or Cloud mode for data persistence, so I start a chroma server here.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@flink.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] [api][runtime] Introduce long-term memory in python [flink-agents]
Posted by "wenjin272 (via GitHub)" <gi...@apache.org>.
wenjin272 commented on code in PR #332:
URL: https://github.com/apache/flink-agents/pull/332#discussion_r2606199703
##########
python/flink_agents/api/vector_stores/vector_store.py:
##########
@@ -85,11 +91,18 @@ class Document(BaseModel):
"""
content: str = Field(description="The actual text content of the document.")
- metadata: Dict[str, Any] = Field(default_factory=dict, description="Document metadata such as source, author, timestamp, etc.")
- id: str | None = Field(default=None, description="Unique identifier of the document.")
+ metadata: Dict[str, Any] = Field(
+ default_factory=dict,
+ description="Document metadata such as source, author, timestamp, etc.",
+ )
+ id: str | None = Field(
+ default=None, description="Unique identifier of the document."
+ )
Review Comment:
I just run `ruff format` before commit the codes. Looks like `ruff check` only checks the lint rule, while `ruff format` only format the code style.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@flink.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] [api][runtime] Introduce long-term memory in python [flink-agents]
Posted by "xintongsong (via GitHub)" <gi...@apache.org>.
xintongsong commented on code in PR #332:
URL: https://github.com/apache/flink-agents/pull/332#discussion_r2605064694
##########
python/flink_agents/api/vector_stores/vector_store.py:
##########
@@ -176,3 +250,74 @@ def query_embedding(self, embedding: List[float], limit: int = 10, **kwargs: Any
Returns:
List of documents matching the search criteria
"""
+
+ @abstractmethod
+ def add_embedding(
+ self,
+ *,
+ documents: List[Document],
+ embeddings: List[List[float]],
+ collection_name: str | None = None,
+ **kwargs: Any,
+ ) -> List[str]:
+ """Add documents with pre-computed embeddings to the vector store.
+
+ Args:
+ documents: Documents to add to the vector store
+ embeddings: Pre-computed embedding vector for each document
+ collection_name: The collection name of the documents to add. Optional.
+ **kwargs: Vector store-specific parameters (collection, namespace, etc.)
+
+ Returns:
+ List of document IDs that were added to the vector store
+ """
+
+
+class Collection(BaseModel):
+ """Represents a collection of documents."""
+ name: str
+ size: int
Review Comment:
What is this size used for and how do we keep it consistent?
##########
python/flink_agents/api/memory/long_term_memory.py:
##########
@@ -0,0 +1,294 @@
+################################################################################
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# 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.
+#################################################################################
+import importlib
+from abc import ABC, abstractmethod
+from datetime import datetime
+from enum import Enum
+from typing import Any, Dict, List, Type
+
+from pydantic import (
+ BaseModel,
+ Field,
+ field_serializer,
+ model_validator,
+)
+from typing_extensions import override
+
+from flink_agents.api.chat_message import ChatMessage
+from flink_agents.api.prompts.prompt import Prompt
+
+ItemType = str | ChatMessage
+
+
+class CompactionStrategyType(Enum):
+ """Strategy for compact memory set."""
+
+ SUMMARIZATION = "summarization"
+
+
+class CompactionStrategy(BaseModel, ABC):
+ """Strategy for compact memory set."""
+
+ @property
+ @abstractmethod
+ def type(self) -> CompactionStrategyType:
+ """Return type of this strategy."""
+
+
+class SummarizationStrategy(CompactionStrategy):
+ """Summarization strategy."""
+
+ model: str
+ prompt: str | Prompt | None = None
+
+ @property
+ @override
+ def type(self) -> CompactionStrategyType:
+ return CompactionStrategyType.SUMMARIZATION
+
+
+class LongTermMemoryBackend(Enum):
+ """Backend for Long-Term Memory."""
+
+ VectorStore = "vectorstore"
Review Comment:
```suggestion
EXTERNAL_VECTOR_STORE = "external_vector_store"
```
##########
python/flink_agents/api/memory/long_term_memory.py:
##########
@@ -0,0 +1,294 @@
+################################################################################
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# 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.
+#################################################################################
+import importlib
+from abc import ABC, abstractmethod
+from datetime import datetime
+from enum import Enum
+from typing import Any, Dict, List, Type
+
+from pydantic import (
+ BaseModel,
+ Field,
+ field_serializer,
+ model_validator,
+)
+from typing_extensions import override
+
+from flink_agents.api.chat_message import ChatMessage
+from flink_agents.api.prompts.prompt import Prompt
+
+ItemType = str | ChatMessage
+
+
+class CompactionStrategyType(Enum):
+ """Strategy for compact memory set."""
+
+ SUMMARIZATION = "summarization"
+
+
+class CompactionStrategy(BaseModel, ABC):
+ """Strategy for compact memory set."""
+
+ @property
+ @abstractmethod
+ def type(self) -> CompactionStrategyType:
+ """Return type of this strategy."""
+
+
+class SummarizationStrategy(CompactionStrategy):
+ """Summarization strategy."""
+
+ model: str
+ prompt: str | Prompt | None = None
+
+ @property
+ @override
+ def type(self) -> CompactionStrategyType:
+ return CompactionStrategyType.SUMMARIZATION
+
+
+class LongTermMemoryBackend(Enum):
+ """Backend for Long-Term Memory."""
+
+ VectorStore = "vectorstore"
+
+
+class DatetimeRange(BaseModel):
+ """Represents a datetime range."""
+
+ start: datetime
+ end: datetime
+
+
+class MemorySetItem(BaseModel):
+ """Represents a long term memory item retrieved from vector store.
+
+ Attributes:
+ memory_set_name: The name of the memory set this item belongs to.
+ id: The id of this item.
+ value: The value of this item.
+ compacted: Whether this item has been compacted.
+ created_time: The timestamp this item was added to the memory set.
+ last_accessed_time: The timestamp this item was last accessed.
+ additional_metadata: Additional metadata for this item.
+ """
+
+ memory_set_name: str
+ id: str
+ value: Any
+ compacted: bool = False
+ created_time: datetime | DatetimeRange = None
+ last_accessed_time: datetime
+ additional_metadata: Dict[str, Any] | None = None
+
+
+class MemorySet(BaseModel):
+ """Represents a long term memory set contains memory items.
+
+ Attributes:
+ name: The name of this memory set.
+ item_type: The type of items stored in this set.
+ capacity: The capacity of this memory set.
+ compaction_strategy: Compaction strategy and additional arguments used
+ to compact memory set.
+ size: The size of this memory set.
+ """
+
+ name: str
+ item_type: Type[str] | Type[ChatMessage]
+ capacity: int
+ compaction_strategy: CompactionStrategy
+ size: int = Field(default=0, exclude=True)
+ ltm: "BaseLongTermMemory" = Field(default=None, exclude=True)
+
+ @field_serializer("item_type")
+ def _serialize_item_type(self, item_type: Type) -> Dict[str, str]:
+ return {"module": item_type.__module__, "name": item_type.__name__}
+
+ @field_serializer("compaction_strategy")
+ def _serialize_compaction_strategy(
+ self, compaction_strategy: CompactionStrategy
+ ) -> Dict[str, str]:
+ data = compaction_strategy.model_dump()
+ data.update(
+ {
+ "module": compaction_strategy.__class__.__module__,
+ "name": compaction_strategy.__class__.__name__,
+ }
+ )
+ return data
+
+ @model_validator(mode="before")
+ def _deserialize_item_type(self) -> "MemorySet":
+ if isinstance(self["item_type"], Dict):
+ module = importlib.import_module(self["item_type"]["module"])
+ self["item_type"] = getattr(module, self["item_type"]["name"])
+ if isinstance(self["compaction_strategy"], Dict):
+ module = importlib.import_module(self["compaction_strategy"].pop("module"))
+ clazz = getattr(module, self["compaction_strategy"].pop("name"))
+ self["compaction_strategy"] = clazz.model_validate(
+ self["compaction_strategy"]
+ )
+ return self
+
+ def add(
+ self, items: ItemType | List[ItemType], ids: str | List[str] | None = None
+ ) -> None:
+ """Add a memory item to the set, currently only support item with
+ type str or ChatMessage.
+
+ If the capacity of this memory set is reached, will trigger reduce
+ operation to manage the memory set size.
+
+ Args:
+ items: The items to be inserted to this set.
+ ids: The ids of the items to be inserted. Optional.
+ """
+ self.ltm.add(memory_set=self, memory_items=items, ids=ids)
+
+ def get(
+ self, ids: str | List[str] | None = None
+ ) -> MemorySetItem | List[MemorySetItem]:
+ """Retrieve memory items. If no item id provided, will return all items.
+
+ Args:
+ ids: The ids of the items to retrieve.
+
+ Returns:
+ The memory items retrieved.
+ """
+ return self.ltm.get(memory_set=self, ids=ids)
+
+ def search(self, query: str, limit: int, **kwargs: Any) -> List[MemorySetItem]:
+ """Retrieve n memory items related to the query.
+
+ Args:
+ query: The query to search for.
+ limit: The number of items to retrieve.
+ **kwargs: Additional arguments for search.
+ """
+ return self.ltm.search(memory_set=self, query=query, limit=limit, **kwargs)
+
+
+class BaseLongTermMemory(ABC, BaseModel):
+ """Base Abstract class for long term memory."""
+
+ @abstractmethod
+ def get_or_create_memory_set(
+ self,
+ name: str,
+ item_type: str | Type[ChatMessage],
+ capacity: int,
+ compaction_strategy: CompactionStrategy,
+ ) -> MemorySet:
+ """Create a memory set, if the memory set already exists, return it.
+
+ Args:
+ name: The name of the memory set.
+ item_type: The type of the memory item.
+ capacity: The capacity of the memory set.
+ compaction_strategy: The compaction strategy and arguments for
+ storge management.
+
+ Returns:
+ The created memory set.
+ """
+
+ @abstractmethod
+ def get_memory_set(self, name: str) -> MemorySet:
+ """Get the memory set.
+
+ Args:
+ name: The name of the memory set.
+
+ Returns:
+ The memory set.
+ """
+
+ @abstractmethod
+ def delete_memory_set(self, name: str) -> MemorySet:
Review Comment:
It probably make sense to just return a boolean. The memory set should be no longer bond with the underlying store.
##########
python/flink_agents/api/vector_stores/vector_store.py:
##########
@@ -176,3 +250,74 @@ def query_embedding(self, embedding: List[float], limit: int = 10, **kwargs: Any
Returns:
List of documents matching the search criteria
"""
+
+ @abstractmethod
+ def add_embedding(
Review Comment:
`add_embedding` and `query_embedding` are not meant for users to call. We probably should mark them as protected.
##########
python/flink_agents/runtime/memory/compaction_functions.py:
##########
@@ -0,0 +1,149 @@
+################################################################################
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# 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 typing import TYPE_CHECKING, List, Type, cast
+
+from flink_agents.api.chat_message import ChatMessage, MessageRole
+from flink_agents.api.memory.long_term_memory import (
+ BaseLongTermMemory,
+ MemorySet,
+ MemorySetItem,
+ SummarizationStrategy,
+)
+from flink_agents.api.resource import ResourceType
+from flink_agents.api.runner_context import RunnerContext
+
+if TYPE_CHECKING:
+ from flink_agents.api.chat_models.chat_model import BaseChatModelSetup
+ from flink_agents.api.prompts.prompt import Prompt
+
+
+def summarize(
Review Comment:
It might require more efforts than we can take in this PR to tune the compaction strategy in order to get an ideal performance. But at least in the abstraction we should not assume this only returns one message.
##########
python/flink_agents/api/memory/long_term_memory.py:
##########
@@ -0,0 +1,294 @@
+################################################################################
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# 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.
+#################################################################################
+import importlib
+from abc import ABC, abstractmethod
+from datetime import datetime
+from enum import Enum
+from typing import Any, Dict, List, Type
+
+from pydantic import (
+ BaseModel,
+ Field,
+ field_serializer,
+ model_validator,
+)
+from typing_extensions import override
+
+from flink_agents.api.chat_message import ChatMessage
+from flink_agents.api.prompts.prompt import Prompt
+
+ItemType = str | ChatMessage
+
+
+class CompactionStrategyType(Enum):
+ """Strategy for compact memory set."""
+
+ SUMMARIZATION = "summarization"
+
+
+class CompactionStrategy(BaseModel, ABC):
+ """Strategy for compact memory set."""
+
+ @property
+ @abstractmethod
+ def type(self) -> CompactionStrategyType:
+ """Return type of this strategy."""
+
+
+class SummarizationStrategy(CompactionStrategy):
+ """Summarization strategy."""
+
+ model: str
+ prompt: str | Prompt | None = None
+
+ @property
+ @override
+ def type(self) -> CompactionStrategyType:
+ return CompactionStrategyType.SUMMARIZATION
+
+
+class LongTermMemoryBackend(Enum):
+ """Backend for Long-Term Memory."""
+
+ VectorStore = "vectorstore"
+
+
+class DatetimeRange(BaseModel):
+ """Represents a datetime range."""
+
+ start: datetime
+ end: datetime
+
+
+class MemorySetItem(BaseModel):
+ """Represents a long term memory item retrieved from vector store.
+
+ Attributes:
+ memory_set_name: The name of the memory set this item belongs to.
+ id: The id of this item.
+ value: The value of this item.
+ compacted: Whether this item has been compacted.
+ created_time: The timestamp this item was added to the memory set.
+ last_accessed_time: The timestamp this item was last accessed.
+ additional_metadata: Additional metadata for this item.
+ """
+
+ memory_set_name: str
+ id: str
+ value: Any
+ compacted: bool = False
+ created_time: datetime | DatetimeRange = None
+ last_accessed_time: datetime
+ additional_metadata: Dict[str, Any] | None = None
+
+
+class MemorySet(BaseModel):
+ """Represents a long term memory set contains memory items.
+
+ Attributes:
+ name: The name of this memory set.
+ item_type: The type of items stored in this set.
+ capacity: The capacity of this memory set.
+ compaction_strategy: Compaction strategy and additional arguments used
+ to compact memory set.
+ size: The size of this memory set.
+ """
+
+ name: str
+ item_type: Type[str] | Type[ChatMessage]
+ capacity: int
+ compaction_strategy: CompactionStrategy
+ size: int = Field(default=0, exclude=True)
Review Comment:
It's fragile to maintain the size separately from the underlying actual collection. We might consider get the size from the store.
##########
python/flink_agents/runtime/memory/compaction_functions.py:
##########
@@ -0,0 +1,149 @@
+################################################################################
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# 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 typing import TYPE_CHECKING, List, Type, cast
+
+from flink_agents.api.chat_message import ChatMessage, MessageRole
+from flink_agents.api.memory.long_term_memory import (
+ BaseLongTermMemory,
+ MemorySet,
+ MemorySetItem,
+ SummarizationStrategy,
+)
+from flink_agents.api.resource import ResourceType
+from flink_agents.api.runner_context import RunnerContext
+
+if TYPE_CHECKING:
+ from flink_agents.api.chat_models.chat_model import BaseChatModelSetup
+ from flink_agents.api.prompts.prompt import Prompt
+
+
+def summarize(
Review Comment:
It doesn't make sense to always compact the entire long term memory into 1 message. I think we should only merge similar messages, and discard meaning less one. As a result, we should get a smaller set of messages with higher information density.
##########
python/flink_agents/api/vector_stores/vector_store.py:
##########
@@ -176,3 +250,74 @@ def query_embedding(self, embedding: List[float], limit: int = 10, **kwargs: Any
Returns:
List of documents matching the search criteria
"""
+
+ @abstractmethod
+ def add_embedding(
+ self,
+ *,
+ documents: List[Document],
+ embeddings: List[List[float]],
+ collection_name: str | None = None,
+ **kwargs: Any,
+ ) -> List[str]:
+ """Add documents with pre-computed embeddings to the vector store.
+
+ Args:
+ documents: Documents to add to the vector store
+ embeddings: Pre-computed embedding vector for each document
+ collection_name: The collection name of the documents to add. Optional.
+ **kwargs: Vector store-specific parameters (collection, namespace, etc.)
+
+ Returns:
+ List of document IDs that were added to the vector store
+ """
+
+
+class Collection(BaseModel):
+ """Represents a collection of documents."""
+ name: str
+ size: int
+ metadata: Dict[str, Any] | None = None
+
+
+class CollectionManageableVectorStore(BaseVectorStore, ABC):
+ """Base abstract class for vector store which support collection management."""
+
+ @abstractmethod
+ def get_or_create_collection(
+ self, name: str, metadata: Dict[str, Any]
+ ) -> Collection:
+ """Get a collection, or create it if it doesn't exist.
+
+ Args:
+ name: Name of the collection
+ metadata: Metadata of the collection
+ Returns:
+ The retrieved or created collection
+ """
+
+ @abstractmethod
+ def get_collection(self, name: str) -> Collection:
+ """Get a collection, raise an exception if it doesn't exist.
+
+ Args:
+ name: Name of the collection
+ Returns:
+ The retrieved collection
+ """
+
+ @abstractmethod
+ def delete_collection(self, name: str) -> Collection:
+ """Delete a collection.
+
+ Args:
+ name: Name of the collection
+ Returns:
+ The deleted collection
+ """
+
+def maybe_cast_to_list(value: Any | List[Any]) -> List[Any] | None:
Review Comment:
Should this be private?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@flink.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] [api][runtime] Introduce long-term memory in python [flink-agents]
Posted by "xintongsong (via GitHub)" <gi...@apache.org>.
xintongsong commented on code in PR #332:
URL: https://github.com/apache/flink-agents/pull/332#discussion_r2573555945
##########
python/flink_agents/api/memory/long_term_memory.py:
##########
@@ -0,0 +1,263 @@
+################################################################################
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# 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.
+#################################################################################
+import importlib
+from abc import ABC, abstractmethod
+from datetime import datetime
+from enum import Enum
+from typing import Any, Dict, List, Type
+
+from pydantic import (
+ BaseModel,
+ Field,
+ field_serializer,
+ model_validator,
+)
+
+from flink_agents.api.chat_message import ChatMessage
+from flink_agents.api.prompts.prompt import Prompt
+
+
+class ReduceStrategy(Enum):
Review Comment:
Not sure about the name `Reduce`. I think `Compact` might be better.
##########
python/flink_agents/api/memory/long_term_memory.py:
##########
@@ -0,0 +1,263 @@
+################################################################################
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# 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.
+#################################################################################
+import importlib
+from abc import ABC, abstractmethod
+from datetime import datetime
+from enum import Enum
+from typing import Any, Dict, List, Type
+
+from pydantic import (
+ BaseModel,
+ Field,
+ field_serializer,
+ model_validator,
+)
+
+from flink_agents.api.chat_message import ChatMessage
+from flink_agents.api.prompts.prompt import Prompt
+
+
+class ReduceStrategy(Enum):
+ """Strategy for reducing memory set size."""
+
+ TRIM = "trim"
+ SUMMARIZE = "summarize"
+
+
+class ReduceSetup(BaseModel):
+ """Reduce setup for managing the storge of memory set."""
+
+ strategy: ReduceStrategy
+ arguments: Dict[str, Any] = Field(default_factory=dict)
+
+ @staticmethod
+ def trim_setup(n: int) -> "ReduceSetup":
+ """Create trim setup."""
+ return ReduceSetup(strategy=ReduceStrategy.TRIM, arguments={"n": n})
+
+ @staticmethod
+ def summarize_setup(
+ n: int, model: str, prompt: Prompt | None = None
+ ) -> "ReduceSetup":
+ """Create summarize setup."""
+ return ReduceSetup(
+ strategy=ReduceStrategy.SUMMARIZE,
+ arguments={"n": n, "model": model, "prompt": prompt},
+ )
+
+
+class LongTermMemoryBackend(Enum):
+ """Backend for Long-Term Memory."""
+
+ CHROMA = "chroma"
+
+
+class DatetimeRange(BaseModel):
+ """Represents a datetime range."""
+
+ start: datetime
+ end: datetime
+
+
+class MemorySetItem(BaseModel):
+ """Represents a long term memory item retrieved from vector store.
+
+ Attributes:
+ memory_set_name: The name of the memory set this item belongs to.
+ id: The id of this item.
+ value: The value of this item.
+ compacted: Whether this item has been compacted.
+ created_time: The timestamp this item was added to the memory set.
+ last_accessed_time: The timestamp this item was last accessed.
+ """
+
+ memory_set_name: str
+ id: str
+ value: Any
+ compacted: bool = False
+ created_time: DatetimeRange
+ last_accessed_time: datetime
+
+class MemorySet(BaseModel):
+ """Represents a long term memory set contains memory items.
+
+ Attributes:
+ name: The name of this memory set.
+ item_type: The type of items stored in this set.
+ size: Current items count stored in this set.
+ capacity: The capacity of this memory set.
+ reduce_setup: Reduce strategy and additional arguments used to reduce memory
+ set size.
+ item_ids: The indices of items stored in this set.
+ reduced: Whether this memory set has been reduced.
+ """
+
+ name: str
+ item_type: Type[str] | Type[ChatMessage]
+ size: int = 0
+ capacity: int
+ reduce_setup: ReduceSetup
+ item_ids: List[str] = Field(default_factory=list)
Review Comment:
Why do we need to store ids of all the items?
##########
python/flink_agents/runtime/memory/chroma_long_term_memory.py:
##########
@@ -0,0 +1,456 @@
+################################################################################
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# 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.
+#################################################################################
+import uuid
+from datetime import datetime, timezone
+from typing import TYPE_CHECKING, Any, Dict, List, Type, cast
+
+import chromadb
+from chromadb import ClientAPI as ChromaClient
+from chromadb import CloudClient, GetResult, Metadata, Settings
+from chromadb.api.types import Document
+from chromadb.errors import NotFoundError
+from pydantic import BaseModel, ConfigDict, Field
+from typing_extensions import override
+
+from flink_agents.api.chat_message import ChatMessage, MessageRole
+from flink_agents.api.embedding_models.embedding_model import BaseEmbeddingModelSetup
+from flink_agents.api.memory.long_term_memory import (
+ BaseLongTermMemory,
+ DatetimeRange,
+ MemorySet,
+ MemorySetItem,
+ ReduceSetup,
+ ReduceStrategy,
+)
+from flink_agents.api.resource import ResourceType
+from flink_agents.api.runner_context import RunnerContext
+from flink_agents.runtime.memory.reduce_functions import summarize
+
+if TYPE_CHECKING:
+ from chromadb import GetResult
+
+
+class ChromaLongTermMemory(BaseLongTermMemory):
+ """Long-Term Memory based on ChromaDB."""
+
+ model_config = ConfigDict(arbitrary_types_allowed=True)
+ embedding_model: str | BaseEmbeddingModelSetup | None = Field(
+ default=None,
+ description="Custom embedding model used to generate item embedding. "
+ "If not provided, will use the default embedding function of chroma.",
+ )
+ runner_context: RunnerContext = Field(
+ default="The runner context for this long term memory.",
+ )
+
+ # Connection configuration
+ persist_directory: str | None = Field(
+ default=None,
+ description="Directory for persistent storage. If None, uses in-memory client.",
+ )
+ host: str | None = Field(
+ default=None,
+ description="Host for ChromaDB server connection.",
+ )
+ port: int | None = Field(
+ default=8000,
+ description="Port for ChromaDB server connection.",
+ )
+ api_key: str | None = Field(
+ default=None,
+ description="API key for Chroma Cloud connection.",
+ )
+ client_settings: Settings | None = Field(
+ default=None,
+ description="ChromaDB client settings for advanced configuration.",
+ )
+ tenant: str = Field(
+ default="default_tenant",
+ description="ChromaDB tenant for multi-tenancy support.",
+ )
+ database: str = Field(
+ default="default_database",
+ description="ChromaDB database name.",
+ )
+
+ __client: ChromaClient | None = None
+
+ def __init__(
+ self,
+ *,
+ runner_context: RunnerContext,
+ embedding_model: str | None = None,
+ persist_directory: str | None = None,
+ host: str | None = None,
+ port: int | None = 8000,
+ api_key: str | None = None,
+ client_settings: Settings | None = None,
+ tenant: str = "default_tenant",
+ database: str,
+ **kwargs: Any,
+ ) -> None:
+ """Init method."""
+ super().__init__(
+ runner_context=runner_context,
+ embedding_model=embedding_model,
+ persist_directory=persist_directory,
+ host=host,
+ port=port,
+ api_key=api_key,
+ client_settings=client_settings,
+ tenant=tenant,
+ database=database,
+ **kwargs,
+ )
+
+ @property
+ def client(self) -> ChromaClient:
+ """Return ChromaDB client, creating it if necessary."""
+ if self.__client is None:
+ # Choose client type based on configuration
+ if self.api_key is not None:
+ # Cloud mode
+ self.__client = CloudClient(
+ tenant=self.tenant,
+ database=self.database,
+ api_key=self.api_key,
+ )
+ elif self.host is not None:
+ # Client-Server Mode
+ # create database if needed
+ admin_client = chromadb.AdminClient(
+ Settings(
+ chroma_api_impl="chromadb.api.fastapi.FastAPI",
+ chroma_server_host=self.host,
+ chroma_server_http_port=self.port,
+ )
+ )
+ try:
+ admin_client.get_database(name=self.database, tenant=self.tenant)
+ except Exception:
+ admin_client.create_database(name=self.database, tenant=self.tenant)
+ self.__client = chromadb.HttpClient(
+ host=self.host,
+ port=self.port,
+ settings=self.client_settings,
+ tenant=self.tenant,
+ database=self.database,
+ )
+ else:
+ err_msg = (
+ "As long-term memory backend, chromadb must run in cloud mode or client-server mode, "
+ "user should provide cloud api-key or server address."
+ )
+ raise RuntimeError(err_msg)
+
+ return self.__client
+
+ @override
+ def create_memory_set(
+ self,
+ name: str,
+ item_type: type[str] | Type[ChatMessage],
+ capacity: int,
+ reduce_setup: ReduceSetup,
+ ) -> MemorySet:
+ try:
+ return self.get_memory_set(name=name)
+ except NotFoundError:
+ memory_set = MemorySet(
+ name=name,
+ item_type=item_type,
+ size=0,
+ capacity=capacity,
+ reduce_setup=reduce_setup,
+ ltm=self,
+ )
+ self.client.get_or_create_collection(
+ name=name, metadata={"serialization": memory_set.model_dump_json()}
+ )
+ return memory_set
+
+ @override
+ def get_memory_set(self, name: str) -> MemorySet:
+ metadata = self.client.get_collection(name=name).metadata
+ memory_set = MemorySet.model_validate_json(metadata["serialization"])
+ memory_set.ltm = self
+ return memory_set
+
+ @override
+ def delete_memory_set(self, name: str) -> None:
+ self.client.delete_collection(name=name)
+
+ @override
+ def add(self, memory_set: MemorySet, memory_item: str | ChatMessage) -> None:
+ if memory_set.size >= memory_set.capacity:
+ # trigger reduce operation to manage memory set size.
+ self._reduce(memory_set)
Review Comment:
This can be extremely slow. We should proactively do the compaction.
##########
python/flink_agents/api/memory/long_term_memory.py:
##########
@@ -0,0 +1,263 @@
+################################################################################
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# 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.
+#################################################################################
+import importlib
+from abc import ABC, abstractmethod
+from datetime import datetime
+from enum import Enum
+from typing import Any, Dict, List, Type
+
+from pydantic import (
+ BaseModel,
+ Field,
+ field_serializer,
+ model_validator,
+)
+
+from flink_agents.api.chat_message import ChatMessage
+from flink_agents.api.prompts.prompt import Prompt
+
+
+class ReduceStrategy(Enum):
+ """Strategy for reducing memory set size."""
+
+ TRIM = "trim"
+ SUMMARIZE = "summarize"
+
+
+class ReduceSetup(BaseModel):
+ """Reduce setup for managing the storge of memory set."""
+
+ strategy: ReduceStrategy
+ arguments: Dict[str, Any] = Field(default_factory=dict)
+
+ @staticmethod
+ def trim_setup(n: int) -> "ReduceSetup":
+ """Create trim setup."""
+ return ReduceSetup(strategy=ReduceStrategy.TRIM, arguments={"n": n})
+
+ @staticmethod
+ def summarize_setup(
+ n: int, model: str, prompt: Prompt | None = None
+ ) -> "ReduceSetup":
+ """Create summarize setup."""
+ return ReduceSetup(
+ strategy=ReduceStrategy.SUMMARIZE,
+ arguments={"n": n, "model": model, "prompt": prompt},
+ )
+
+
+class LongTermMemoryBackend(Enum):
+ """Backend for Long-Term Memory."""
+
+ CHROMA = "chroma"
Review Comment:
We probably should either support a embedded backend only, or support any external vector store.
##########
python/flink_agents/api/memory/long_term_memory.py:
##########
@@ -0,0 +1,263 @@
+################################################################################
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# 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.
+#################################################################################
+import importlib
+from abc import ABC, abstractmethod
+from datetime import datetime
+from enum import Enum
+from typing import Any, Dict, List, Type
+
+from pydantic import (
+ BaseModel,
+ Field,
+ field_serializer,
+ model_validator,
+)
+
+from flink_agents.api.chat_message import ChatMessage
+from flink_agents.api.prompts.prompt import Prompt
+
+
+class ReduceStrategy(Enum):
+ """Strategy for reducing memory set size."""
+
+ TRIM = "trim"
+ SUMMARIZE = "summarize"
+
+
+class ReduceSetup(BaseModel):
+ """Reduce setup for managing the storge of memory set."""
+
+ strategy: ReduceStrategy
+ arguments: Dict[str, Any] = Field(default_factory=dict)
+
+ @staticmethod
+ def trim_setup(n: int) -> "ReduceSetup":
+ """Create trim setup."""
+ return ReduceSetup(strategy=ReduceStrategy.TRIM, arguments={"n": n})
+
+ @staticmethod
+ def summarize_setup(
+ n: int, model: str, prompt: Prompt | None = None
+ ) -> "ReduceSetup":
+ """Create summarize setup."""
+ return ReduceSetup(
+ strategy=ReduceStrategy.SUMMARIZE,
+ arguments={"n": n, "model": model, "prompt": prompt},
+ )
+
+
+class LongTermMemoryBackend(Enum):
+ """Backend for Long-Term Memory."""
+
+ CHROMA = "chroma"
+
+
+class DatetimeRange(BaseModel):
+ """Represents a datetime range."""
+
+ start: datetime
+ end: datetime
+
+
+class MemorySetItem(BaseModel):
+ """Represents a long term memory item retrieved from vector store.
+
+ Attributes:
+ memory_set_name: The name of the memory set this item belongs to.
+ id: The id of this item.
+ value: The value of this item.
+ compacted: Whether this item has been compacted.
+ created_time: The timestamp this item was added to the memory set.
+ last_accessed_time: The timestamp this item was last accessed.
+ """
+
+ memory_set_name: str
+ id: str
+ value: Any
+ compacted: bool = False
+ created_time: DatetimeRange
+ last_accessed_time: datetime
+
+class MemorySet(BaseModel):
+ """Represents a long term memory set contains memory items.
+
+ Attributes:
+ name: The name of this memory set.
+ item_type: The type of items stored in this set.
+ size: Current items count stored in this set.
+ capacity: The capacity of this memory set.
+ reduce_setup: Reduce strategy and additional arguments used to reduce memory
+ set size.
+ item_ids: The indices of items stored in this set.
+ reduced: Whether this memory set has been reduced.
+ """
+
+ name: str
+ item_type: Type[str] | Type[ChatMessage]
+ size: int = 0
+ capacity: int
+ reduce_setup: ReduceSetup
+ item_ids: List[str] = Field(default_factory=list)
+ reduced: bool = False
+ ltm: "BaseLongTermMemory" = Field(default=None, exclude=True)
+
+ @field_serializer("item_type")
+ def _serialize_item_type(self, item_type: Type) -> Dict[str, str]:
+ return {"module": item_type.__module__, "name": item_type.__name__}
+
+ @model_validator(mode="before")
+ def _deserialize_item_type(self) -> "MemorySet":
+ if isinstance(self["item_type"], Dict):
+ module = importlib.import_module(self["item_type"]["module"])
+ self["item_type"] = getattr(module, self["item_type"]["name"])
+ return self
+
+ def add(self, item: str | ChatMessage) -> None:
+ """Add a memory item to the set, currently only support item with
+ type str or ChatMessage.
+
+ If the capacity of this memory set is reached, will trigger reduce
+ operation to manage the memory set size.
+
+ Args:
+ item: The item to be inserted to this set.
+ """
+ self.ltm.add(memory_set=self, memory_item=item)
+
+ def get(self) -> List[MemorySetItem]:
+ """Retrieve all memory items.
+
+ Returns:
+ All memory items in this set.
+ """
+ return self.ltm.get(memory_set=self)
+
+ def get_recent(self, n: int) -> List[MemorySetItem]:
+ """Retrieve n most recent memory items.
+
+ Args:
+ n: The number of items to retrieve.
+
+ Returns:
+ List of memory items retrieved, sorted by creation timestamp.
+ """
+ return self.ltm.get_recent(memory_set=self, n=n)
+
+ def search(self, query: str, limit: int, **kwargs: Any) -> List[MemorySetItem]:
+ """Retrieve n memory items related to the query.
+
+ Args:
+ query: The query to search for.
+ limit: The number of items to retrieve.
+ **kwargs: Additional arguments for search.
+ """
+ return self.ltm.search(memory_set=self, query=query, limit=limit, **kwargs)
+
+
+class BaseLongTermMemory(ABC, BaseModel):
+ """Base Abstract class for long term memory."""
+
+ @abstractmethod
+ def create_memory_set(
+ self,
+ name: str,
+ item_type: str | Type[ChatMessage],
+ capacity: int,
+ reduce_setup: ReduceSetup,
+ ) -> MemorySet:
+ """Create a memory set, if the memory set already exists, return it.
+
+ Args:
+ name: The name of the memory set.
+ item_type: The type of the memory item.
+ capacity: The capacity of the memory set.
+ reduce_setup: The reduce strategy and arguments for storge management.
+
+ Returns:
+ The created memory set.
+ """
+
+ @abstractmethod
+ def get_memory_set(self, name: str) -> MemorySet:
+ """Get the memory set.
+
+ Args:
+ name: The name of the memory set.
+
+ Returns:
+ The memory set.
+ """
+
+ @abstractmethod
+ def delete_memory_set(self, name: str) -> None:
Review Comment:
Usually, for a delete method, we return a boolean of the deleted element / null to indicate whether the element has been deleted.
##########
python/flink_agents/api/memory/long_term_memory.py:
##########
@@ -0,0 +1,263 @@
+################################################################################
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# 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.
+#################################################################################
+import importlib
+from abc import ABC, abstractmethod
+from datetime import datetime
+from enum import Enum
+from typing import Any, Dict, List, Type
+
+from pydantic import (
+ BaseModel,
+ Field,
+ field_serializer,
+ model_validator,
+)
+
+from flink_agents.api.chat_message import ChatMessage
+from flink_agents.api.prompts.prompt import Prompt
+
+
+class ReduceStrategy(Enum):
+ """Strategy for reducing memory set size."""
+
+ TRIM = "trim"
+ SUMMARIZE = "summarize"
+
+
+class ReduceSetup(BaseModel):
+ """Reduce setup for managing the storge of memory set."""
+
+ strategy: ReduceStrategy
+ arguments: Dict[str, Any] = Field(default_factory=dict)
+
+ @staticmethod
+ def trim_setup(n: int) -> "ReduceSetup":
+ """Create trim setup."""
+ return ReduceSetup(strategy=ReduceStrategy.TRIM, arguments={"n": n})
+
+ @staticmethod
+ def summarize_setup(
+ n: int, model: str, prompt: Prompt | None = None
+ ) -> "ReduceSetup":
+ """Create summarize setup."""
+ return ReduceSetup(
+ strategy=ReduceStrategy.SUMMARIZE,
+ arguments={"n": n, "model": model, "prompt": prompt},
+ )
+
+
+class LongTermMemoryBackend(Enum):
+ """Backend for Long-Term Memory."""
+
+ CHROMA = "chroma"
+
+
+class DatetimeRange(BaseModel):
+ """Represents a datetime range."""
+
+ start: datetime
+ end: datetime
+
+
+class MemorySetItem(BaseModel):
+ """Represents a long term memory item retrieved from vector store.
+
+ Attributes:
+ memory_set_name: The name of the memory set this item belongs to.
+ id: The id of this item.
+ value: The value of this item.
+ compacted: Whether this item has been compacted.
+ created_time: The timestamp this item was added to the memory set.
+ last_accessed_time: The timestamp this item was last accessed.
+ """
+
+ memory_set_name: str
+ id: str
+ value: Any
+ compacted: bool = False
+ created_time: DatetimeRange
Review Comment:
```suggestion
created_time: DatetimeRange | datetime
```
##########
python/flink_agents/api/memory/long_term_memory.py:
##########
@@ -0,0 +1,263 @@
+################################################################################
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# 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.
+#################################################################################
+import importlib
+from abc import ABC, abstractmethod
+from datetime import datetime
+from enum import Enum
+from typing import Any, Dict, List, Type
+
+from pydantic import (
+ BaseModel,
+ Field,
+ field_serializer,
+ model_validator,
+)
+
+from flink_agents.api.chat_message import ChatMessage
+from flink_agents.api.prompts.prompt import Prompt
+
+
+class ReduceStrategy(Enum):
+ """Strategy for reducing memory set size."""
+
+ TRIM = "trim"
+ SUMMARIZE = "summarize"
+
+
+class ReduceSetup(BaseModel):
+ """Reduce setup for managing the storge of memory set."""
+
+ strategy: ReduceStrategy
+ arguments: Dict[str, Any] = Field(default_factory=dict)
+
+ @staticmethod
+ def trim_setup(n: int) -> "ReduceSetup":
+ """Create trim setup."""
+ return ReduceSetup(strategy=ReduceStrategy.TRIM, arguments={"n": n})
+
+ @staticmethod
+ def summarize_setup(
+ n: int, model: str, prompt: Prompt | None = None
+ ) -> "ReduceSetup":
+ """Create summarize setup."""
+ return ReduceSetup(
+ strategy=ReduceStrategy.SUMMARIZE,
+ arguments={"n": n, "model": model, "prompt": prompt},
+ )
+
+
+class LongTermMemoryBackend(Enum):
+ """Backend for Long-Term Memory."""
+
+ CHROMA = "chroma"
+
+
+class DatetimeRange(BaseModel):
+ """Represents a datetime range."""
+
+ start: datetime
+ end: datetime
+
+
+class MemorySetItem(BaseModel):
+ """Represents a long term memory item retrieved from vector store.
+
+ Attributes:
+ memory_set_name: The name of the memory set this item belongs to.
+ id: The id of this item.
+ value: The value of this item.
+ compacted: Whether this item has been compacted.
+ created_time: The timestamp this item was added to the memory set.
+ last_accessed_time: The timestamp this item was last accessed.
+ """
+
+ memory_set_name: str
+ id: str
+ value: Any
+ compacted: bool = False
+ created_time: DatetimeRange
+ last_accessed_time: datetime
+
+class MemorySet(BaseModel):
+ """Represents a long term memory set contains memory items.
+
+ Attributes:
+ name: The name of this memory set.
+ item_type: The type of items stored in this set.
+ size: Current items count stored in this set.
+ capacity: The capacity of this memory set.
+ reduce_setup: Reduce strategy and additional arguments used to reduce memory
+ set size.
+ item_ids: The indices of items stored in this set.
+ reduced: Whether this memory set has been reduced.
+ """
+
+ name: str
+ item_type: Type[str] | Type[ChatMessage]
+ size: int = 0
+ capacity: int
+ reduce_setup: ReduceSetup
+ item_ids: List[str] = Field(default_factory=list)
+ reduced: bool = False
Review Comment:
And why do we need to know whether a memory set has been reduced/compacted?
##########
python/flink_agents/api/memory/long_term_memory.py:
##########
@@ -0,0 +1,263 @@
+################################################################################
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# 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.
+#################################################################################
+import importlib
+from abc import ABC, abstractmethod
+from datetime import datetime
+from enum import Enum
+from typing import Any, Dict, List, Type
+
+from pydantic import (
+ BaseModel,
+ Field,
+ field_serializer,
+ model_validator,
+)
+
+from flink_agents.api.chat_message import ChatMessage
+from flink_agents.api.prompts.prompt import Prompt
+
+
+class ReduceStrategy(Enum):
+ """Strategy for reducing memory set size."""
+
+ TRIM = "trim"
+ SUMMARIZE = "summarize"
+
+
+class ReduceSetup(BaseModel):
Review Comment:
I think `CompactionStrategy.trim(n)` might be more straightforward for users, compared to `ReduceSetup.trim_setup(n)`.
##########
python/flink_agents/runtime/memory/chroma_long_term_memory.py:
##########
@@ -0,0 +1,456 @@
+################################################################################
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# 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.
+#################################################################################
+import uuid
+from datetime import datetime, timezone
+from typing import TYPE_CHECKING, Any, Dict, List, Type, cast
+
+import chromadb
+from chromadb import ClientAPI as ChromaClient
+from chromadb import CloudClient, GetResult, Metadata, Settings
+from chromadb.api.types import Document
+from chromadb.errors import NotFoundError
+from pydantic import BaseModel, ConfigDict, Field
+from typing_extensions import override
+
+from flink_agents.api.chat_message import ChatMessage, MessageRole
+from flink_agents.api.embedding_models.embedding_model import BaseEmbeddingModelSetup
+from flink_agents.api.memory.long_term_memory import (
+ BaseLongTermMemory,
+ DatetimeRange,
+ MemorySet,
+ MemorySetItem,
+ ReduceSetup,
+ ReduceStrategy,
+)
+from flink_agents.api.resource import ResourceType
+from flink_agents.api.runner_context import RunnerContext
+from flink_agents.runtime.memory.reduce_functions import summarize
+
+if TYPE_CHECKING:
+ from chromadb import GetResult
+
+
+class ChromaLongTermMemory(BaseLongTermMemory):
+ """Long-Term Memory based on ChromaDB."""
+
+ model_config = ConfigDict(arbitrary_types_allowed=True)
+ embedding_model: str | BaseEmbeddingModelSetup | None = Field(
+ default=None,
+ description="Custom embedding model used to generate item embedding. "
+ "If not provided, will use the default embedding function of chroma.",
+ )
+ runner_context: RunnerContext = Field(
+ default="The runner context for this long term memory.",
+ )
+
+ # Connection configuration
+ persist_directory: str | None = Field(
+ default=None,
+ description="Directory for persistent storage. If None, uses in-memory client.",
+ )
+ host: str | None = Field(
+ default=None,
+ description="Host for ChromaDB server connection.",
+ )
+ port: int | None = Field(
+ default=8000,
+ description="Port for ChromaDB server connection.",
+ )
+ api_key: str | None = Field(
+ default=None,
+ description="API key for Chroma Cloud connection.",
+ )
+ client_settings: Settings | None = Field(
+ default=None,
+ description="ChromaDB client settings for advanced configuration.",
+ )
+ tenant: str = Field(
+ default="default_tenant",
+ description="ChromaDB tenant for multi-tenancy support.",
+ )
+ database: str = Field(
+ default="default_database",
+ description="ChromaDB database name.",
+ )
+
+ __client: ChromaClient | None = None
+
+ def __init__(
+ self,
+ *,
+ runner_context: RunnerContext,
+ embedding_model: str | None = None,
+ persist_directory: str | None = None,
+ host: str | None = None,
+ port: int | None = 8000,
+ api_key: str | None = None,
+ client_settings: Settings | None = None,
+ tenant: str = "default_tenant",
+ database: str,
+ **kwargs: Any,
+ ) -> None:
+ """Init method."""
+ super().__init__(
+ runner_context=runner_context,
+ embedding_model=embedding_model,
+ persist_directory=persist_directory,
+ host=host,
+ port=port,
+ api_key=api_key,
+ client_settings=client_settings,
+ tenant=tenant,
+ database=database,
+ **kwargs,
+ )
+
+ @property
+ def client(self) -> ChromaClient:
+ """Return ChromaDB client, creating it if necessary."""
+ if self.__client is None:
+ # Choose client type based on configuration
+ if self.api_key is not None:
+ # Cloud mode
+ self.__client = CloudClient(
+ tenant=self.tenant,
+ database=self.database,
+ api_key=self.api_key,
+ )
+ elif self.host is not None:
+ # Client-Server Mode
+ # create database if needed
+ admin_client = chromadb.AdminClient(
+ Settings(
+ chroma_api_impl="chromadb.api.fastapi.FastAPI",
+ chroma_server_host=self.host,
+ chroma_server_http_port=self.port,
+ )
+ )
+ try:
+ admin_client.get_database(name=self.database, tenant=self.tenant)
+ except Exception:
+ admin_client.create_database(name=self.database, tenant=self.tenant)
+ self.__client = chromadb.HttpClient(
+ host=self.host,
+ port=self.port,
+ settings=self.client_settings,
+ tenant=self.tenant,
+ database=self.database,
+ )
+ else:
+ err_msg = (
+ "As long-term memory backend, chromadb must run in cloud mode or client-server mode, "
+ "user should provide cloud api-key or server address."
+ )
+ raise RuntimeError(err_msg)
+
+ return self.__client
+
+ @override
+ def create_memory_set(
+ self,
+ name: str,
+ item_type: type[str] | Type[ChatMessage],
+ capacity: int,
+ reduce_setup: ReduceSetup,
+ ) -> MemorySet:
+ try:
+ return self.get_memory_set(name=name)
+ except NotFoundError:
+ memory_set = MemorySet(
+ name=name,
+ item_type=item_type,
+ size=0,
+ capacity=capacity,
+ reduce_setup=reduce_setup,
+ ltm=self,
+ )
+ self.client.get_or_create_collection(
+ name=name, metadata={"serialization": memory_set.model_dump_json()}
+ )
+ return memory_set
+
+ @override
+ def get_memory_set(self, name: str) -> MemorySet:
+ metadata = self.client.get_collection(name=name).metadata
+ memory_set = MemorySet.model_validate_json(metadata["serialization"])
+ memory_set.ltm = self
+ return memory_set
+
+ @override
+ def delete_memory_set(self, name: str) -> None:
+ self.client.delete_collection(name=name)
+
+ @override
+ def add(self, memory_set: MemorySet, memory_item: str | ChatMessage) -> None:
+ if memory_set.size >= memory_set.capacity:
+ # trigger reduce operation to manage memory set size.
+ self._reduce(memory_set)
+
+ if issubclass(memory_set.item_type, BaseModel):
+ memory_item = memory_item.model_dump_json()
+
+ embedding = self._generate_embedding(text=memory_item)
+
+ item_id = str(uuid.uuid4())
+ timestamp = datetime.now(timezone.utc).isoformat()
+ self.client.get_collection(name=memory_set.name).add(
+ ids=item_id,
+ embeddings=embedding,
+ documents=memory_item,
+ metadatas={
+ "compacted": False,
+ "created_time_start": timestamp,
+ "created_time_end": timestamp,
+ "last_accessed_time": timestamp,
+ },
+ )
+ memory_set.size = memory_set.size + 1
+ memory_set.item_ids.append(item_id)
+ self._update_set_metadata(memory_set)
+
+ @override
+ def get(self, memory_set: MemorySet) -> List[MemorySetItem]:
+ return self.slice(memory_set=memory_set)
+
+ @override
+ def get_recent(self, memory_set: MemorySet, n: int) -> List[MemorySetItem]:
+ offset = memory_set.size - n if memory_set.size > n else 0
+ return self.slice(memory_set=memory_set, offset=offset, n=n)
+
+ @override
+ def search(
Review Comment:
Same here.
##########
python/flink_agents/runtime/memory/chroma_long_term_memory.py:
##########
@@ -0,0 +1,456 @@
+################################################################################
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# 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.
+#################################################################################
+import uuid
+from datetime import datetime, timezone
+from typing import TYPE_CHECKING, Any, Dict, List, Type, cast
+
+import chromadb
+from chromadb import ClientAPI as ChromaClient
+from chromadb import CloudClient, GetResult, Metadata, Settings
+from chromadb.api.types import Document
+from chromadb.errors import NotFoundError
+from pydantic import BaseModel, ConfigDict, Field
+from typing_extensions import override
+
+from flink_agents.api.chat_message import ChatMessage, MessageRole
+from flink_agents.api.embedding_models.embedding_model import BaseEmbeddingModelSetup
+from flink_agents.api.memory.long_term_memory import (
+ BaseLongTermMemory,
+ DatetimeRange,
+ MemorySet,
+ MemorySetItem,
+ ReduceSetup,
+ ReduceStrategy,
+)
+from flink_agents.api.resource import ResourceType
+from flink_agents.api.runner_context import RunnerContext
+from flink_agents.runtime.memory.reduce_functions import summarize
+
+if TYPE_CHECKING:
+ from chromadb import GetResult
+
+
+class ChromaLongTermMemory(BaseLongTermMemory):
+ """Long-Term Memory based on ChromaDB."""
+
+ model_config = ConfigDict(arbitrary_types_allowed=True)
+ embedding_model: str | BaseEmbeddingModelSetup | None = Field(
+ default=None,
+ description="Custom embedding model used to generate item embedding. "
+ "If not provided, will use the default embedding function of chroma.",
+ )
+ runner_context: RunnerContext = Field(
+ default="The runner context for this long term memory.",
+ )
+
+ # Connection configuration
+ persist_directory: str | None = Field(
+ default=None,
+ description="Directory for persistent storage. If None, uses in-memory client.",
+ )
+ host: str | None = Field(
+ default=None,
+ description="Host for ChromaDB server connection.",
+ )
+ port: int | None = Field(
+ default=8000,
+ description="Port for ChromaDB server connection.",
+ )
+ api_key: str | None = Field(
+ default=None,
+ description="API key for Chroma Cloud connection.",
+ )
+ client_settings: Settings | None = Field(
+ default=None,
+ description="ChromaDB client settings for advanced configuration.",
+ )
+ tenant: str = Field(
+ default="default_tenant",
+ description="ChromaDB tenant for multi-tenancy support.",
+ )
+ database: str = Field(
+ default="default_database",
+ description="ChromaDB database name.",
+ )
+
+ __client: ChromaClient | None = None
+
+ def __init__(
+ self,
+ *,
+ runner_context: RunnerContext,
+ embedding_model: str | None = None,
+ persist_directory: str | None = None,
+ host: str | None = None,
+ port: int | None = 8000,
+ api_key: str | None = None,
+ client_settings: Settings | None = None,
+ tenant: str = "default_tenant",
+ database: str,
+ **kwargs: Any,
+ ) -> None:
+ """Init method."""
+ super().__init__(
+ runner_context=runner_context,
+ embedding_model=embedding_model,
+ persist_directory=persist_directory,
+ host=host,
+ port=port,
+ api_key=api_key,
+ client_settings=client_settings,
+ tenant=tenant,
+ database=database,
+ **kwargs,
+ )
+
+ @property
+ def client(self) -> ChromaClient:
+ """Return ChromaDB client, creating it if necessary."""
+ if self.__client is None:
+ # Choose client type based on configuration
+ if self.api_key is not None:
+ # Cloud mode
+ self.__client = CloudClient(
+ tenant=self.tenant,
+ database=self.database,
+ api_key=self.api_key,
+ )
+ elif self.host is not None:
+ # Client-Server Mode
+ # create database if needed
+ admin_client = chromadb.AdminClient(
+ Settings(
+ chroma_api_impl="chromadb.api.fastapi.FastAPI",
+ chroma_server_host=self.host,
+ chroma_server_http_port=self.port,
+ )
+ )
+ try:
+ admin_client.get_database(name=self.database, tenant=self.tenant)
+ except Exception:
+ admin_client.create_database(name=self.database, tenant=self.tenant)
+ self.__client = chromadb.HttpClient(
+ host=self.host,
+ port=self.port,
+ settings=self.client_settings,
+ tenant=self.tenant,
+ database=self.database,
+ )
+ else:
+ err_msg = (
+ "As long-term memory backend, chromadb must run in cloud mode or client-server mode, "
+ "user should provide cloud api-key or server address."
+ )
+ raise RuntimeError(err_msg)
+
+ return self.__client
+
+ @override
+ def create_memory_set(
+ self,
+ name: str,
+ item_type: type[str] | Type[ChatMessage],
+ capacity: int,
+ reduce_setup: ReduceSetup,
+ ) -> MemorySet:
+ try:
+ return self.get_memory_set(name=name)
+ except NotFoundError:
+ memory_set = MemorySet(
+ name=name,
+ item_type=item_type,
+ size=0,
+ capacity=capacity,
+ reduce_setup=reduce_setup,
+ ltm=self,
+ )
+ self.client.get_or_create_collection(
+ name=name, metadata={"serialization": memory_set.model_dump_json()}
+ )
+ return memory_set
+
+ @override
+ def get_memory_set(self, name: str) -> MemorySet:
+ metadata = self.client.get_collection(name=name).metadata
+ memory_set = MemorySet.model_validate_json(metadata["serialization"])
+ memory_set.ltm = self
+ return memory_set
+
+ @override
+ def delete_memory_set(self, name: str) -> None:
+ self.client.delete_collection(name=name)
+
+ @override
+ def add(self, memory_set: MemorySet, memory_item: str | ChatMessage) -> None:
Review Comment:
I had a feeling that adding items to long-term memory can take time, for embedding. We probably should also provide async apis.
##########
python/flink_agents/api/memory/long_term_memory.py:
##########
@@ -0,0 +1,263 @@
+################################################################################
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# 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.
+#################################################################################
+import importlib
+from abc import ABC, abstractmethod
+from datetime import datetime
+from enum import Enum
+from typing import Any, Dict, List, Type
+
+from pydantic import (
+ BaseModel,
+ Field,
+ field_serializer,
+ model_validator,
+)
+
+from flink_agents.api.chat_message import ChatMessage
+from flink_agents.api.prompts.prompt import Prompt
+
+
+class ReduceStrategy(Enum):
+ """Strategy for reducing memory set size."""
+
+ TRIM = "trim"
+ SUMMARIZE = "summarize"
+
+
+class ReduceSetup(BaseModel):
Review Comment:
I'd suggest to name this `CompactionStrategy`, and make it an abstract class that we can provide different implementations, so we can have strict limit on which arguments should be specified for each strategy. We can call the current `ReduceStrategy` `CompactionStrategyType`.
##########
python/flink_agents/runtime/memory/chroma_long_term_memory.py:
##########
@@ -0,0 +1,456 @@
+################################################################################
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# 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.
+#################################################################################
+import uuid
+from datetime import datetime, timezone
+from typing import TYPE_CHECKING, Any, Dict, List, Type, cast
+
+import chromadb
+from chromadb import ClientAPI as ChromaClient
+from chromadb import CloudClient, GetResult, Metadata, Settings
+from chromadb.api.types import Document
+from chromadb.errors import NotFoundError
+from pydantic import BaseModel, ConfigDict, Field
+from typing_extensions import override
+
+from flink_agents.api.chat_message import ChatMessage, MessageRole
+from flink_agents.api.embedding_models.embedding_model import BaseEmbeddingModelSetup
+from flink_agents.api.memory.long_term_memory import (
+ BaseLongTermMemory,
+ DatetimeRange,
+ MemorySet,
+ MemorySetItem,
+ ReduceSetup,
+ ReduceStrategy,
+)
+from flink_agents.api.resource import ResourceType
+from flink_agents.api.runner_context import RunnerContext
+from flink_agents.runtime.memory.reduce_functions import summarize
+
+if TYPE_CHECKING:
+ from chromadb import GetResult
+
+
+class ChromaLongTermMemory(BaseLongTermMemory):
+ """Long-Term Memory based on ChromaDB."""
+
+ model_config = ConfigDict(arbitrary_types_allowed=True)
+ embedding_model: str | BaseEmbeddingModelSetup | None = Field(
+ default=None,
+ description="Custom embedding model used to generate item embedding. "
+ "If not provided, will use the default embedding function of chroma.",
+ )
+ runner_context: RunnerContext = Field(
+ default="The runner context for this long term memory.",
+ )
+
+ # Connection configuration
+ persist_directory: str | None = Field(
Review Comment:
what is this directory for?
##########
python/flink_agents/runtime/memory/chroma_long_term_memory.py:
##########
@@ -0,0 +1,456 @@
+################################################################################
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# 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.
+#################################################################################
+import uuid
+from datetime import datetime, timezone
+from typing import TYPE_CHECKING, Any, Dict, List, Type, cast
+
+import chromadb
+from chromadb import ClientAPI as ChromaClient
+from chromadb import CloudClient, GetResult, Metadata, Settings
+from chromadb.api.types import Document
+from chromadb.errors import NotFoundError
+from pydantic import BaseModel, ConfigDict, Field
+from typing_extensions import override
+
+from flink_agents.api.chat_message import ChatMessage, MessageRole
+from flink_agents.api.embedding_models.embedding_model import BaseEmbeddingModelSetup
+from flink_agents.api.memory.long_term_memory import (
+ BaseLongTermMemory,
+ DatetimeRange,
+ MemorySet,
+ MemorySetItem,
+ ReduceSetup,
+ ReduceStrategy,
+)
+from flink_agents.api.resource import ResourceType
+from flink_agents.api.runner_context import RunnerContext
+from flink_agents.runtime.memory.reduce_functions import summarize
+
+if TYPE_CHECKING:
+ from chromadb import GetResult
+
+
+class ChromaLongTermMemory(BaseLongTermMemory):
+ """Long-Term Memory based on ChromaDB."""
+
+ model_config = ConfigDict(arbitrary_types_allowed=True)
+ embedding_model: str | BaseEmbeddingModelSetup | None = Field(
+ default=None,
+ description="Custom embedding model used to generate item embedding. "
+ "If not provided, will use the default embedding function of chroma.",
+ )
+ runner_context: RunnerContext = Field(
+ default="The runner context for this long term memory.",
+ )
+
+ # Connection configuration
+ persist_directory: str | None = Field(
+ default=None,
+ description="Directory for persistent storage. If None, uses in-memory client.",
+ )
+ host: str | None = Field(
+ default=None,
+ description="Host for ChromaDB server connection.",
+ )
+ port: int | None = Field(
+ default=8000,
+ description="Port for ChromaDB server connection.",
+ )
+ api_key: str | None = Field(
+ default=None,
+ description="API key for Chroma Cloud connection.",
+ )
+ client_settings: Settings | None = Field(
+ default=None,
+ description="ChromaDB client settings for advanced configuration.",
+ )
+ tenant: str = Field(
+ default="default_tenant",
+ description="ChromaDB tenant for multi-tenancy support.",
+ )
+ database: str = Field(
+ default="default_database",
+ description="ChromaDB database name.",
+ )
+
+ __client: ChromaClient | None = None
+
+ def __init__(
+ self,
+ *,
+ runner_context: RunnerContext,
+ embedding_model: str | None = None,
+ persist_directory: str | None = None,
+ host: str | None = None,
+ port: int | None = 8000,
+ api_key: str | None = None,
+ client_settings: Settings | None = None,
+ tenant: str = "default_tenant",
+ database: str,
+ **kwargs: Any,
+ ) -> None:
+ """Init method."""
+ super().__init__(
+ runner_context=runner_context,
+ embedding_model=embedding_model,
+ persist_directory=persist_directory,
+ host=host,
+ port=port,
+ api_key=api_key,
+ client_settings=client_settings,
+ tenant=tenant,
+ database=database,
+ **kwargs,
+ )
+
+ @property
+ def client(self) -> ChromaClient:
+ """Return ChromaDB client, creating it if necessary."""
+ if self.__client is None:
+ # Choose client type based on configuration
+ if self.api_key is not None:
+ # Cloud mode
+ self.__client = CloudClient(
+ tenant=self.tenant,
+ database=self.database,
+ api_key=self.api_key,
+ )
+ elif self.host is not None:
+ # Client-Server Mode
+ # create database if needed
+ admin_client = chromadb.AdminClient(
+ Settings(
+ chroma_api_impl="chromadb.api.fastapi.FastAPI",
+ chroma_server_host=self.host,
+ chroma_server_http_port=self.port,
+ )
+ )
+ try:
+ admin_client.get_database(name=self.database, tenant=self.tenant)
+ except Exception:
+ admin_client.create_database(name=self.database, tenant=self.tenant)
+ self.__client = chromadb.HttpClient(
+ host=self.host,
+ port=self.port,
+ settings=self.client_settings,
+ tenant=self.tenant,
+ database=self.database,
+ )
+ else:
+ err_msg = (
+ "As long-term memory backend, chromadb must run in cloud mode or client-server mode, "
+ "user should provide cloud api-key or server address."
+ )
+ raise RuntimeError(err_msg)
+
+ return self.__client
+
+ @override
+ def create_memory_set(
+ self,
+ name: str,
+ item_type: type[str] | Type[ChatMessage],
+ capacity: int,
+ reduce_setup: ReduceSetup,
+ ) -> MemorySet:
+ try:
+ return self.get_memory_set(name=name)
+ except NotFoundError:
+ memory_set = MemorySet(
+ name=name,
+ item_type=item_type,
+ size=0,
+ capacity=capacity,
+ reduce_setup=reduce_setup,
+ ltm=self,
+ )
+ self.client.get_or_create_collection(
+ name=name, metadata={"serialization": memory_set.model_dump_json()}
+ )
+ return memory_set
+
+ @override
+ def get_memory_set(self, name: str) -> MemorySet:
+ metadata = self.client.get_collection(name=name).metadata
+ memory_set = MemorySet.model_validate_json(metadata["serialization"])
+ memory_set.ltm = self
+ return memory_set
+
+ @override
+ def delete_memory_set(self, name: str) -> None:
+ self.client.delete_collection(name=name)
+
+ @override
+ def add(self, memory_set: MemorySet, memory_item: str | ChatMessage) -> None:
+ if memory_set.size >= memory_set.capacity:
+ # trigger reduce operation to manage memory set size.
+ self._reduce(memory_set)
+
+ if issubclass(memory_set.item_type, BaseModel):
+ memory_item = memory_item.model_dump_json()
+
+ embedding = self._generate_embedding(text=memory_item)
+
+ item_id = str(uuid.uuid4())
+ timestamp = datetime.now(timezone.utc).isoformat()
+ self.client.get_collection(name=memory_set.name).add(
+ ids=item_id,
+ embeddings=embedding,
+ documents=memory_item,
+ metadatas={
+ "compacted": False,
+ "created_time_start": timestamp,
+ "created_time_end": timestamp,
+ "last_accessed_time": timestamp,
+ },
+ )
+ memory_set.size = memory_set.size + 1
+ memory_set.item_ids.append(item_id)
+ self._update_set_metadata(memory_set)
+
+ @override
+ def get(self, memory_set: MemorySet) -> List[MemorySetItem]:
+ return self.slice(memory_set=memory_set)
+
+ @override
+ def get_recent(self, memory_set: MemorySet, n: int) -> List[MemorySetItem]:
+ offset = memory_set.size - n if memory_set.size > n else 0
+ return self.slice(memory_set=memory_set, offset=offset, n=n)
+
+ @override
+ def search(
+ self, memory_set: MemorySet, query: str, limit: int, **kwargs: Any
+ ) -> List[MemorySetItem]:
+ embedding = self._generate_embedding(text=query)
+ result = self.client.get_collection(name=memory_set.name).query(
+ query_embeddings=[embedding] if embedding else None,
+ query_texts=[query],
+ n_results=limit,
+ where=kwargs.get("where"),
+ include=["documents", "metadatas"],
+ )
+
+ ids = result["ids"][0]
+ documents = result["documents"][0]
+ metadatas = result["metadatas"][0]
+
+ self._update_items_metadata(memory_set=memory_set, ids=ids, metadatas=metadatas)
+
+ return self._convert_to_items(
+ memory_set=memory_set, ids=ids, documents=documents, metadatas=metadatas
+ )
+
+ def slice(
+ self,
+ *,
+ memory_set: MemorySet,
+ offset: int | None = None,
+ n: int | None = None,
+ update_metadata: bool = True,
+ ) -> List[MemorySetItem]:
+ """Retrieve memory items up to limit starting at offset.
+
+ Args:
+ memory_set: The memory set to be retrieved from.
+ offset: The offset to start retrieving from.
+ n: The number of items to retrieve.
+ update_metadata: If True, update the items metadata in chroma store.
+
+ Returns:
+ Retrieved memory items.
+ """
+ result: GetResult = self.client.get_collection(name=memory_set.name).get(
+ offset=offset, limit=n
+ )
+
+ ids = result["ids"]
+ metadatas = result["metadatas"]
+ documents = result["documents"]
+
+ if update_metadata:
+ self._update_items_metadata(
+ memory_set=memory_set, ids=ids, metadatas=metadatas
+ )
+
+ return self._convert_to_items(
+ memory_set=memory_set, ids=ids, documents=documents, metadatas=metadatas
+ )
+
+ def delete(self, memory_set: MemorySet, offset: int, n: int) -> None:
+ """Delete memory items.
+
+ Args:
+ memory_set: The memory set to be deleted from.
+ offset: The offset to start delete from.
+ n: The number of items to delete.
+ """
+ self.client.get_collection(name=memory_set.name).delete(
+ memory_set.item_ids[offset : offset + n]
+ )
+ del memory_set.item_ids[offset : offset + n]
+ memory_set.size = memory_set.size - n
+
+ self.client.get_collection(name=memory_set.name).modify(
+ metadata={"serialization": memory_set.model_dump_json()}
+ )
+
+ def update(
+ self, memory_set: MemorySet, item_id: str, text: str, metadata: Dict[str, Any]
+ ) -> None:
+ """Update memory item.
+
+ Args:
+ memory_set: The memory set item belongs to.
+ item_id: The memory item to be updated.
+ text: The updated text.
+ metadata: The updated metadata.
+ """
+ embedding = self._generate_embedding(text=text)
+ self.client.get_collection(name=memory_set.name).update(
+ ids=item_id,
+ embeddings=embedding,
+ documents=text,
+ metadatas=metadata,
+ )
+
+ @staticmethod
+ def _convert_to_items(
+ memory_set: MemorySet,
+ ids: List[str],
+ documents: List[Document],
+ metadatas: List[Metadata],
+ ) -> List[MemorySetItem]:
+ """Convert retrival result to memory items."""
+ return [
+ MemorySetItem(
+ memory_set_name=memory_set.name,
+ id=item_id,
+ value=document
+ if memory_set.item_type is str
+ else memory_set.item_type.model_validate_json(document),
+ compacted=metadata["compacted"],
+ created_time=DatetimeRange(
+ start=datetime.fromisoformat(
+ cast("str", metadata["created_time_start"])
+ ),
+ end=datetime.fromisoformat(cast("str", metadata["created_time_end"])),
+ ),
+ last_accessed_time=datetime.fromisoformat(metadata["last_accessed_time"]),
+ )
+ for item_id, document, metadata in zip(
+ ids, documents, metadatas, strict=False
+ )
+ ]
+
+ def _generate_embedding(self, text: str) -> list[float] | None:
+ """Generate embedding for text.
+
+ If no embedding model configured, return None to use the default embedding
+ generated by chroma.
+ """
+ if self.embedding_model is not None:
+ if isinstance(self.embedding_model, str):
+ self.embedding_model = cast(
+ "BaseEmbeddingModelSetup",
+ self.runner_context.get_resource(
+ self.embedding_model, ResourceType.EMBEDDING_MODEL
+ ),
+ )
+ return self.embedding_model.embed(text=text)
+ else:
+ return None
+
+ def _update_set_metadata(self, memory_set: MemorySet) -> None:
+ """Update metadata for memory set."""
+ self.client.get_collection(name=memory_set.name).modify(
+ metadata={"serialization": memory_set.model_dump_json()}
+ )
+
+ def _update_items_metadata(
+ self, memory_set: MemorySet, ids: List[str], metadatas: List[Metadata]
+ ) -> None:
+ """Update metadata for retrieved memory items."""
+ current_timestamp = datetime.now(timezone.utc).isoformat()
+ update_metadatas = []
+ for metadata in metadatas:
+ metadata = dict(metadata)
+ metadata["last_accessed_time"] = current_timestamp
+ update_metadatas.append(metadata)
+ self.client.get_collection(name=memory_set.name).update(
+ ids=ids, metadatas=update_metadatas
+ )
+
+ def _reduce(self, memory_set: MemorySet) -> None:
+ """Reduce memory set size."""
+ reduce_setup: ReduceSetup = memory_set.reduce_setup
+ if reduce_setup.strategy == ReduceStrategy.TRIM:
+ self._trim(memory_set)
+ elif reduce_setup.strategy == ReduceStrategy.SUMMARIZE:
+ self._summarize(memory_set)
+ else:
+ msg = f"Unknown reduce strategy: {reduce_setup.strategy}"
+ raise RuntimeError(msg)
+ memory_set.reduced = True
+ self._update_set_metadata(memory_set)
+
+ def _trim(self, memory_set: MemorySet) -> None:
+ reduce_setup: ReduceSetup = memory_set.reduce_setup
+ n = reduce_setup.arguments.get("n")
+ self.delete(memory_set=memory_set, offset=0, n=n)
+
+ def _summarize(self, memory_set: MemorySet) -> None:
Review Comment:
Are these methods specialized for this class?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@flink.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] [api][runtime] Introduce long-term memory in python [flink-agents]
Posted by "xintongsong (via GitHub)" <gi...@apache.org>.
xintongsong merged PR #332:
URL: https://github.com/apache/flink-agents/pull/332
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@flink.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] [api][runtime] Introduce long-term memory in python [flink-agents]
Posted by "alnzng (via GitHub)" <gi...@apache.org>.
alnzng commented on code in PR #332:
URL: https://github.com/apache/flink-agents/pull/332#discussion_r2605447094
##########
python/flink_agents/api/vector_stores/vector_store.py:
##########
@@ -85,11 +91,18 @@ class Document(BaseModel):
"""
content: str = Field(description="The actual text content of the document.")
- metadata: Dict[str, Any] = Field(default_factory=dict, description="Document metadata such as source, author, timestamp, etc.")
- id: str | None = Field(default=None, description="Unique identifier of the document.")
+ metadata: Dict[str, Any] = Field(
+ default_factory=dict,
+ description="Document metadata such as source, author, timestamp, etc.",
+ )
+ id: str | None = Field(
+ default=None, description="Unique identifier of the document."
+ )
Review Comment:
Out of curiosity, did we change the code format stype? seems there are many diffs like this
##########
python/flink_agents/api/vector_stores/vector_store.py:
##########
@@ -135,6 +152,40 @@ def store_kwargs(self) -> Dict[str, Any]:
when performing vector search operations.
"""
+ def add(
+ self, documents: Document | List[Document], collection_name: str | None = None, **kwargs: Any
Review Comment:
The newly added CRUD methods look flexible enough to handle operations across multiple collections, which is good. However, the current vector store implementation is limited to a single collection. To keep the behavior consistent, we should update the existing implementation to support multiple collections as well. For example, For example, the `query`implementation needs to be extended to support retrieval from multiple collections to align with the new CRUD methods.
##########
python/flink_agents/api/vector_stores/vector_store.py:
##########
@@ -135,6 +152,40 @@ def store_kwargs(self) -> Dict[str, Any]:
when performing vector search operations.
"""
+ def add(
+ self, documents: Document | List[Document], collection_name: str | None = None, **kwargs: Any
+ ) -> List[str]:
+ """Add documents to the vector store.
+
+ Converts document content to embeddings and stores them in the vector store.
+ The implementation may generate IDs for documents that don't have one.
+
+ Args:
+ documents: Single Document or list of Documents to add.
+ collection_name: The collection name of the documents to add to. Optional.
+ **kwargs: Vector store specific parameters.
+
+ Returns:
+ List of document IDs that were added to the vector store
+ """
+ # Normalize to list
+ documents = maybe_cast_to_list(documents)
+
+ # Generate embeddings for all documents
+ embedding_model = self.get_resource(
+ self.embedding_model, ResourceType.EMBEDDING_MODEL
+ )
+
+ # Generate embeddings for each document
+ embeddings = [embedding_model.embed(doc.content) for doc in documents]
Review Comment:
Chroma seems to have a size limitation on each request when persisting embeddings, batches should contain fewer than 41,666 embeddings. This has been discussed in the Chroma community (see [GitHub issue #1049](https://github.com/chroma-core/chroma/issues/1049)). I’m not sure if this restriction has changed recently, but to be safe we should add a check for requests exceeding 41,666 embeddings, similar to how LlamaIndex handles it today: https://github.com/run-llama/llama_index/blob/main/llama-index-integrations/vector_stores/llama-index-vector-stores-chroma/llama_index/vector_stores/chroma/base.py#L97C1-L97C72
That being said, it may be better to abstract the embedding function at the document level and allow each vector store implementation to override it as needed.
##########
python/flink_agents/runtime/memory/tests/start_chroma_server.sh:
##########
@@ -0,0 +1,23 @@
+#!/usr/bin/env bash
+
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# 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.
+#
+#
+
+path=$1
+chroma run --path $path
Review Comment:
To enable PersistentClient[1], specify `persist_directory` parameter while building the `ChromaVectorStore` should help. IIUC, this is similar to what "--path $path" parameter offer in chroma CLI.
[1] https://docs.trychroma.com/docs/run-chroma/persistent-client
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@flink.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] [api][runtime] Introduce long-term memory in python [flink-agents]
Posted by "wenjin272 (via GitHub)" <gi...@apache.org>.
wenjin272 commented on PR #332:
URL: https://github.com/apache/flink-agents/pull/332#issuecomment-3595429303
Hi, @alnzng. There's a design issue related to the vector store that I'd appreciate your help reviewing.
As describe in the design doc #339, long-term memory of flink-agents is also based on vector store. Currently, I provide an implementation based on chroma. In this implementation, I directly use chroma client rather than flink-agents `BaseVectorStore`, because there are some long-term memory needed interface not provided in `BaseVectorStore`.
@xintongsong believes that we can directly build upon the Flink-Agents `BaseVectorStore`. Thus, we can support using any already supported vector store as the backend for long-term memory.
I think this make sense, but it requires add some interfaces to `BaseVectorStore`, which maybe look like:
```
def get_or_create_collection(self, name: str, metadata: Dict[str, Any]) -> Dict[str, Any]:
"""Get a collection, create if it doesn't already exist."""
def get_collection(self, name: str) -> None:
"""Get an existing collection."""
def update_collection(self, name: str, metadata: Dict[str, Any]) -> None:
"""Update an existing collection."""
def delete_collection(self, name: str) -> bool:
"""Delete a collection."""
def add(self, document: Document, collection_name: str | None = None) -> None:
"""Add a document to the collection."""
def update(self, document: Document, collection_name: str | None = None) -> None:
"""Update a document, can only update metadata."""
def delete(self, offset: int | None, limit: int | None, ids: List[int] | None = None, **kwargs: Any) -> bool:
"""Delete documents from collection."""
def get(self, offset: int | None, limit: int | None, collection_name: str | None = None, **kwargs: Any) -> List[Document]:
"""Get documents from collection."""
```
WDTY?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@flink.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] [api][runtime] Introduce long-term memory in python [flink-agents]
Posted by "alnzng (via GitHub)" <gi...@apache.org>.
alnzng commented on code in PR #332:
URL: https://github.com/apache/flink-agents/pull/332#discussion_r2583199262
##########
python/flink_agents/runtime/memory/tests/start_chroma_server.sh:
##########
@@ -0,0 +1,23 @@
+#!/usr/bin/env bash
+
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# 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.
+#
+#
+
+path=$1
+chroma run --path $path
Review Comment:
is this script must to have?
IIUC, You can just add "import chromadb" in the test class which will "start" the chroma automatically. I did this in `test_chroma_vector_store`.
##########
python/flink_agents/api/memory/long_term_memory.py:
##########
@@ -0,0 +1,263 @@
+################################################################################
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# 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.
+#################################################################################
+import importlib
+from abc import ABC, abstractmethod
+from datetime import datetime
+from enum import Enum
+from typing import Any, Dict, List, Type
+
+from pydantic import (
+ BaseModel,
+ Field,
+ field_serializer,
+ model_validator,
+)
+
+from flink_agents.api.chat_message import ChatMessage
+from flink_agents.api.prompts.prompt import Prompt
+
+
+class ReduceStrategy(Enum):
+ """Strategy for reducing memory set size."""
+
+ TRIM = "trim"
+ SUMMARIZE = "summarize"
+
+
+class ReduceSetup(BaseModel):
Review Comment:
- +1 on that `Compaction` is a better terminology. This `Compaction` term is commonly used in many open source software in the industry today.
- Should we consider using noun as the name of each strategy. For example, `summarize` -> `summarization`, `trim` -> `truncation`?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@flink.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] [api][runtime] Introduce long-term memory in python [flink-agents]
Posted by "wenjin272 (via GitHub)" <gi...@apache.org>.
wenjin272 commented on PR #332:
URL: https://github.com/apache/flink-agents/pull/332#issuecomment-3604787435
> > Hi, @alnzng. There's a design issue related to the vector store that I'd appreciate your help reviewing.
> > As describe in the design doc #339, long-term memory of flink-agents is also based on vector store. Currently, I provide an implementation based on chroma. In this implementation, I directly use chroma client rather than flink-agents `BaseVectorStore`, because there are some long-term memory needed interface not provided in `BaseVectorStore`.
> > @xintongsong believes that we can directly build upon the Flink-Agents `BaseVectorStore`. Thus, we can support using any already supported vector store as the backend for long-term memory.
> > I think this make sense, but it requires add some interfaces to `BaseVectorStore`, which maybe look like:
> > ```
> > def get_or_create_collection(self, name: str, metadata: Dict[str, Any]) -> Dict[str, Any]:
> > """Get a collection, create if it doesn't already exist."""
> >
> > def get_collection(self, name: str) -> None:
> > """Get an existing collection."""
> >
> > def update_collection(self, name: str, metadata: Dict[str, Any]) -> None:
> > """Update an existing collection."""
> >
> > def delete_collection(self, name: str) -> bool:
> > """Delete a collection."""
> >
> > def add(self, document: Document, collection_name: str | None = None) -> None:
> > """Add a document to the collection."""
> >
> > def update(self, document: Document, collection_name: str | None = None) -> None:
> > """Update a document, can only update metadata."""
> >
> > def delete(self, offset: int | None, limit: int | None, ids: List[int] | None = None, **kwargs: Any) -> bool:
> > """Delete documents from collection."""
> >
> > def get(self, offset: int | None, limit: int | None, collection_name: str | None = None, **kwargs: Any) -> List[Document]:
> > """Get documents from collection."""
> > ```
> >
> >
> >
> >
> >
> >
> >
> >
> >
> >
> >
> > These interface may not be achievable for each vector store, I will conduct research and refinement afterward. WDTY?
>
> Thanks for raising this design question @wenjin272! I agree with @xintongsong that building long-term memory on top of `BaseVectorStore` would be beneficial for consistency and reusability.
>
> The current `BaseVectorStore` is intentionally minimal, focusing on read-only semantic search. This design works well for RAG retrieval use cases but indeed lacks the CRUD operations needed for long-term memory management. Adding document-level CRUD operations to `BaseVectorStore` makes sense, as these operations are generally supported across all major vector stores.
>
> A few thoughts on the proposed methods:
>
> * add(), delete(), get() - Widely supported, should be included
> * update() - I'm not certain about common use cases for updating existing long-term memory entries, but it shouldn't be
> problematic to include.
>
> However, I have concerns about including collection management in the base interface. This concept is vector store specific and may cause integration issues. "Collection" is most common term (e.g. Chroma, Milvus, Qdrant, etc), but Pinecone and Weaviate use different terms(e.g. Class, Namespace, etc).
>
> My suggestion: Keep collection management implementation specific rather than in the base interface. Each vector store integration can expose collection management using its native terminology and patterns.
>
> This approach aligns with how `LlamaIndex` and `LangChain` handle vector stores. I'd recommend you reviewing these implementations, they provide good precedents for balancing abstraction with flexibility.
>
> * LlamaIndex BasePydanticVectorStore:
> https://github.com/run-llama/llama_index/blob/main/llama-index-core/llama_index/core/vector_stores/types.py#L334
> * LangChain VectorStore: https://github.com/langchain-ai/langchain/blob/master/libs/core/langchain_core/vectorstores/base.py#L43
Thanks for your suggestion @alnzng.
The reason I want to provide collection level operation in `BaseVectorStore` is `LongTermMemory` can dynamic create or delete `MemorySet`, which is correspond to a `Collection` in vector store.
I have reviewed the vector store design in langChain, it indeed doesn't provide collection management in base interface. But if we keep collection management implementation specific rather than in the base interface, the `LongTermMemory` must be aware of what vector store implementation is used when create or get a `MemorySet`.
I think all the vector store provide concept like `collection`, they may called `index`, `namespace` or `class`. The collection in `BaseVectorStore` is represent a flink-agents collection, and can be correspond to `collection` in milvus, chroma, quota , `index` in opensearchpy and `class` or `namespace` in pinecone or weaviate. Because collection is the most used name, so we use collection also.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@flink.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] [api][runtime] Introduce long-term memory in python [flink-agents]
Posted by "wenjin272 (via GitHub)" <gi...@apache.org>.
wenjin272 commented on PR #332:
URL: https://github.com/apache/flink-agents/pull/332#issuecomment-3631750287
Hi, @alnzng, I add some interfaces to `BaseVectorStore` and add a new base class `CollectionManageableVectorStore` for vector store could provide collection level management. Please take a look at your convenience, thanks.
The async execution for `add`, `search`, `get`, `delete` and compaction for long-term memory will be implemented in follow PRs.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@flink.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] [api][runtime] Introduce long-term memory in python [flink-agents]
Posted by "alnzng (via GitHub)" <gi...@apache.org>.
alnzng commented on PR #332:
URL: https://github.com/apache/flink-agents/pull/332#issuecomment-3605320883
>I think all the vector store provide concept like collection, they may called index, namespace or class.
@wenjin272 This was my major concern part since I was not able to verify all existing vector stores and we can't predict for new future vector stores. For example, Pinecone uses namespace + index to manage its data, namespace which seems like a "Collection". However seems it doesn't seem to support namespace(medata) update operation, which means below method won't work with Pinecone.
```
def update_collection(self, name: str, metadata: Dict[str, Any]) -> None:
```
@xintongsong suggestion on adding a new interface for indicating collection support seems to be a safe solution.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@flink.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] [api][runtime] Introduce long-term memory in python [flink-agents]
Posted by "xintongsong (via GitHub)" <gi...@apache.org>.
xintongsong commented on code in PR #332:
URL: https://github.com/apache/flink-agents/pull/332#discussion_r2621906683
##########
python/flink_agents/runtime/memory/compaction_functions.py:
##########
@@ -0,0 +1,173 @@
+################################################################################
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# 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.
+#################################################################################
+import json
+from typing import TYPE_CHECKING, List, Type, cast
+
+from flink_agents.api.chat_message import ChatMessage, MessageRole
+from flink_agents.api.memory.long_term_memory import (
+ BaseLongTermMemory,
+ MemorySet,
+ MemorySetItem,
+ SummarizationStrategy,
+)
+from flink_agents.api.prompts.prompt import Prompt
+from flink_agents.api.resource import ResourceType
+from flink_agents.api.runner_context import RunnerContext
+
+if TYPE_CHECKING:
+ from flink_agents.api.chat_models.chat_model import BaseChatModelSetup
+
+
+DEFAULT_ANALYSIS_PROMPT = Prompt.from_text("""<role>
+Context Summarize Assistant
+</role>
+
+<primary_objective>
+Your sole objective in this task is to summarize the context above.
+</primary_objective>
+
+<objective_information>
+You're nearing the total number of input tokens you can accept, so you need compact the context. To achieve this objective, you must extract important topics first. The extracted topics
+must no more than {limit}. Afterwards, you should generate summarization for each topic, and and record which messages the summary was derived from.
+</objective_information>
+
+<output_example>
+You must always respond with valid json format in this format:
+{"topic1": {"summarization": "User ask what is 1 * 2, and the result is 3.", "messages": [0,1,2]},
+ ...
+ "topic4": {"summarization": "User ask what's the weather tomorrow, llm use the search_weather, and the answer is snow.", "messages": [9,10,11,12]}
+}
+</output_example>
+""")
+
+
+def summarize(
+ ltm: BaseLongTermMemory,
+ memory_set: MemorySet,
+ ctx: RunnerContext,
+ ids: List[str] | None = None,
+) -> None:
+ """Generate summarization of the items in the memory set.
+
+ Will add the summarization to memory set, and delete original items involved
+ in summarization.
+
+ Args:
+ ltm: The long term memory the memory set belongs to.
+ memory_set: The memory set to be summarized.
+ ctx: The runner context used to retrieve needed resources.
+ ids: The ids of items to be summarized. If not provided, all items will be
+ involved in summarization. Optional
+ """
+ strategy: SummarizationStrategy = cast(
+ "SummarizationStrategy", memory_set.compaction_strategy
+ )
+
+ # retrieve all items
+ items: List[MemorySetItem] = ltm.get(memory_set=memory_set, ids=ids)
+
+ response: ChatMessage = _generate_summarization(
+ items, memory_set.item_type, strategy, ctx
+ )
Review Comment:
The amount of memory to be compacted can be large and may exceed the context window of the model. We should take that into consideration. Let's add a todo here and create a follow-up issue.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@flink.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org