From 87248586126be17adff4f567ae80569ea1f2fff9 Mon Sep 17 00:00:00 2001 From: "fangyaozheng@bytedance.com" Date: Wed, 31 Dec 2025 09:17:51 +0800 Subject: [PATCH 1/2] feat: support langchain and langgraph community --- .../langchain_ai/checkpoint/__init__.py | 13 +++ .../checkpoint/memory/__init__.py | 13 +++ .../langchain_ai/middlewares/__init__.py | 13 +++ .../langchain_ai/middlewares/save_session.py | 60 +++++++++++ .../community/langchain_ai/models/__init__.py | 13 +++ .../langchain_ai/models/ark_model.py | 34 ++++++ .../community/langchain_ai/store/__init__.py | 13 +++ .../langchain_ai/store/memory/__init__.py | 13 +++ .../store/memory/viking_memory.py | 102 ++++++++++++++++++ .../community/langchain_ai/tools/__init__.py | 13 +++ .../langchain_ai/tools/load_knowledgebase.py | 34 ++++++ .../langchain_ai/tools/load_memory.py | 43 ++++++++ .../vikingdb_memory_backend.py | 11 +- 13 files changed, 372 insertions(+), 3 deletions(-) create mode 100644 veadk/community/langchain_ai/checkpoint/__init__.py create mode 100644 veadk/community/langchain_ai/checkpoint/memory/__init__.py create mode 100644 veadk/community/langchain_ai/middlewares/__init__.py create mode 100644 veadk/community/langchain_ai/middlewares/save_session.py create mode 100644 veadk/community/langchain_ai/models/__init__.py create mode 100644 veadk/community/langchain_ai/models/ark_model.py create mode 100644 veadk/community/langchain_ai/store/__init__.py create mode 100644 veadk/community/langchain_ai/store/memory/__init__.py create mode 100644 veadk/community/langchain_ai/store/memory/viking_memory.py create mode 100644 veadk/community/langchain_ai/tools/__init__.py create mode 100644 veadk/community/langchain_ai/tools/load_knowledgebase.py create mode 100644 veadk/community/langchain_ai/tools/load_memory.py diff --git a/veadk/community/langchain_ai/checkpoint/__init__.py b/veadk/community/langchain_ai/checkpoint/__init__.py new file mode 100644 index 00000000..7f463206 --- /dev/null +++ b/veadk/community/langchain_ai/checkpoint/__init__.py @@ -0,0 +1,13 @@ +# Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates. +# +# Licensed 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. diff --git a/veadk/community/langchain_ai/checkpoint/memory/__init__.py b/veadk/community/langchain_ai/checkpoint/memory/__init__.py new file mode 100644 index 00000000..7f463206 --- /dev/null +++ b/veadk/community/langchain_ai/checkpoint/memory/__init__.py @@ -0,0 +1,13 @@ +# Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates. +# +# Licensed 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. diff --git a/veadk/community/langchain_ai/middlewares/__init__.py b/veadk/community/langchain_ai/middlewares/__init__.py new file mode 100644 index 00000000..7f463206 --- /dev/null +++ b/veadk/community/langchain_ai/middlewares/__init__.py @@ -0,0 +1,13 @@ +# Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates. +# +# Licensed 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. diff --git a/veadk/community/langchain_ai/middlewares/save_session.py b/veadk/community/langchain_ai/middlewares/save_session.py new file mode 100644 index 00000000..d1b99186 --- /dev/null +++ b/veadk/community/langchain_ai/middlewares/save_session.py @@ -0,0 +1,60 @@ +# Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates. +# +# Licensed 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 __future__ import annotations + +from langchain.agents import AgentState +from langchain.agents.middleware import after_agent +from langchain_core.messages.ai import AIMessage +from langchain_core.messages.human import HumanMessage +from langgraph.runtime import Runtime + +from veadk.community.langchain_ai.store.memory.viking_memory import ( + VikingMemoryStore, +) +from veadk.utils.logger import get_logger + +logger = get_logger(__name__) + + +@after_agent +def save_session(state: AgentState, runtime: Runtime) -> None: + """Save the session to the memory store.""" + store: VikingMemoryStore | None = runtime.store + if not store: + return + + app_name = store.index + user_id = runtime.context.user_id + session_id = runtime.context.session_id + + messages = state.get("messages", []) + logger.debug( + f"Save session {session_id} for user {user_id} with {len(messages)} messages. messages={messages}" + ) + + events = {} + for message in messages: + print(type(message)) + if isinstance(message, HumanMessage): + event = {"role": "user", "parts": [{"text": message.content}]} + + elif isinstance(message, AIMessage): + event = {"role": "assistant", "parts": [{"text": message.content}]} + else: + ... + + events[message.id] = event + + store.put(namespace=(app_name, user_id), key=session_id, value=events) diff --git a/veadk/community/langchain_ai/models/__init__.py b/veadk/community/langchain_ai/models/__init__.py new file mode 100644 index 00000000..7f463206 --- /dev/null +++ b/veadk/community/langchain_ai/models/__init__.py @@ -0,0 +1,13 @@ +# Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates. +# +# Licensed 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. diff --git a/veadk/community/langchain_ai/models/ark_model.py b/veadk/community/langchain_ai/models/ark_model.py new file mode 100644 index 00000000..8e48bf3d --- /dev/null +++ b/veadk/community/langchain_ai/models/ark_model.py @@ -0,0 +1,34 @@ +# Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates. +# +# Licensed 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 __future__ import annotations + +from langchain_core.utils.utils import ( + from_env, + secret_from_env, +) +from langchain_openai import ChatOpenAI + + +class ArkChatModel(ChatOpenAI): + def __init__(self, model: str, **kwargs): + super().__init__( + model=model, + api_key=secret_from_env("MODEL_AGENT_API_KEY")(), + base_url=from_env( + "MODEL_AGENT_API_BASE_URL", + default="https://ark.cn-beijing.volces.com/api/v3", + )(), + **kwargs, + ) diff --git a/veadk/community/langchain_ai/store/__init__.py b/veadk/community/langchain_ai/store/__init__.py new file mode 100644 index 00000000..7f463206 --- /dev/null +++ b/veadk/community/langchain_ai/store/__init__.py @@ -0,0 +1,13 @@ +# Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates. +# +# Licensed 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. diff --git a/veadk/community/langchain_ai/store/memory/__init__.py b/veadk/community/langchain_ai/store/memory/__init__.py new file mode 100644 index 00000000..7f463206 --- /dev/null +++ b/veadk/community/langchain_ai/store/memory/__init__.py @@ -0,0 +1,13 @@ +# Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates. +# +# Licensed 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. diff --git a/veadk/community/langchain_ai/store/memory/viking_memory.py b/veadk/community/langchain_ai/store/memory/viking_memory.py new file mode 100644 index 00000000..018d15e5 --- /dev/null +++ b/veadk/community/langchain_ai/store/memory/viking_memory.py @@ -0,0 +1,102 @@ +# Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates. +# +# Licensed 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 __future__ import annotations + +import json +from collections.abc import Iterable + +from langgraph.store.base import ( + BaseStore, + GetOp, + ListNamespacesOp, + Op, + PutOp, + Result, + SearchOp, +) + +from veadk.memory.long_term_memory_backends.vikingdb_memory_backend import ( + VikingDBLTMBackend, +) + + +# mem0 +# viking db +class VikingMemoryStore(BaseStore): + def __init__(self, index: str): + # index = (index, user_id) + # key = session_id + self.index = index + self._backend = VikingDBLTMBackend(index=index) + + def batch(self, ops: Iterable[Op]) -> list[Result]: + # The batch/abatch methods are treated as internal. + # Users should access via put/search/get/list_namespaces/etc. + results = [] + for op in ops: + if isinstance(op, PutOp): + self._apply_put_op(op) + elif isinstance(op, GetOp): + self._apply_get_op(op) + elif isinstance(op, SearchOp): + results.extend(self._apply_search_op(op)) + # elif isinstance(op, ListNamespacesOp): + # self._apply_list_namespaces_op(op) + else: + raise ValueError(f"Unknown op type: {type(op)}") + + return results + + def abatch( + self, ops: Iterable[GetOp | SearchOp | PutOp | ListNamespacesOp] + ) -> list[Result]: ... + + def _apply_put_op(self, op: PutOp) -> None: + index, user_id = op.namespace + session_id = op.key + + assert index == self._backend.index, ( + "index must be the same as the backend index" + ) + + value = op.value + + event_strings = [] + + for _, event in value.items(): + event_strings.append(json.dumps(event)) + + if self._backend.save_memory( + user_id=user_id, + session_id=session_id, + event_strings=event_strings, + ): + return None + + def _apply_get_op(self, op: GetOp): + return ["Not implemented"] + + def _apply_search_op(self, op: SearchOp): + index, user_id = op.namespace_prefix + assert index == self._backend.index, ( + "index must be the same as the backend index" + ) + + query = op.query + if not query: + return [] + + value = self._backend.search_memory(user_id=user_id, query=query, top_k=1) + return value diff --git a/veadk/community/langchain_ai/tools/__init__.py b/veadk/community/langchain_ai/tools/__init__.py new file mode 100644 index 00000000..7f463206 --- /dev/null +++ b/veadk/community/langchain_ai/tools/__init__.py @@ -0,0 +1,13 @@ +# Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates. +# +# Licensed 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. diff --git a/veadk/community/langchain_ai/tools/load_knowledgebase.py b/veadk/community/langchain_ai/tools/load_knowledgebase.py new file mode 100644 index 00000000..3a38b8c7 --- /dev/null +++ b/veadk/community/langchain_ai/tools/load_knowledgebase.py @@ -0,0 +1,34 @@ +# Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates. +# +# Licensed 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 langchain.tools import ToolRuntime, tool + +from veadk.knowledgebase import KnowledgeBase +from veadk.utils.logger import get_logger + +logger = get_logger(__name__) + + +@tool +def load_knowledgebase(query: str, runtime: ToolRuntime) -> list[str]: + """Load knowledge base for the current user. + + Args: + query: The query to search for in the knowledge base. + """ + knowledgeabse: KnowledgeBase = runtime.context.knowledgebase # type: ignore + + results = knowledgeabse.search(query) + + return [result.content for result in results] diff --git a/veadk/community/langchain_ai/tools/load_memory.py b/veadk/community/langchain_ai/tools/load_memory.py new file mode 100644 index 00000000..65b5432c --- /dev/null +++ b/veadk/community/langchain_ai/tools/load_memory.py @@ -0,0 +1,43 @@ +# Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates. +# +# Licensed 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 langchain.tools import ToolRuntime, tool + +from veadk.community.langchain_ai.store.memory.viking_memory import ( + VikingMemoryStore, +) +from veadk.utils.logger import get_logger + +logger = get_logger(__name__) + + +@tool +def load_memory(query: str, runtime: ToolRuntime) -> list[str]: + """Load memories for the current user across all history sessions. + + Args: + query: The query to search for in the memory. + """ + store: VikingMemoryStore | None = runtime.store # type: ignore + if not store: + return ["Long-term memory store is not initialized."] + + app_name = store.index + user_id = runtime.context.user_id + + logger.info(f"Load memory for user {user_id} with query {query}") + response = store.search((app_name, user_id), query=query) + + return response diff --git a/veadk/memory/long_term_memory_backends/vikingdb_memory_backend.py b/veadk/memory/long_term_memory_backends/vikingdb_memory_backend.py index b5f905da..361b1772 100644 --- a/veadk/memory/long_term_memory_backends/vikingdb_memory_backend.py +++ b/veadk/memory/long_term_memory_backends/vikingdb_memory_backend.py @@ -143,9 +143,14 @@ def _get_sdk_client(self) -> VikingMem: ) @override - def save_memory(self, user_id: str, event_strings: list[str], **kwargs) -> bool: + def save_memory( + self, + user_id: str, + event_strings: list[str], + **kwargs, + ) -> bool: assistant_id = kwargs.get("assistant_id", "assistant") - session_id = str(uuid.uuid1()) + session_id = kwargs.get("session_id", str(uuid.uuid1())) messages = [] for raw_events in event_strings: event = json.loads(raw_events) @@ -161,7 +166,7 @@ def save_memory(self, user_id: str, event_strings: list[str], **kwargs) -> bool: } logger.debug( - f"Request for add {len(messages)} memory to VikingDB: collection_name={self.index}, metadata={metadata}, session_id={session_id}" + f"Request for add {len(messages)} memory to VikingDB: collection_name={self.index}, metadata={metadata}, session_id={session_id}, messages={messages}" ) client = self._get_sdk_client() From 472a74692d78ce9743e8adf1da5acc8349dfaa54 Mon Sep 17 00:00:00 2001 From: "fangyaozheng@bytedance.com" Date: Wed, 31 Dec 2025 11:28:35 +0800 Subject: [PATCH 2/2] fix base url name --- veadk/community/langchain_ai/models/ark_model.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/veadk/community/langchain_ai/models/ark_model.py b/veadk/community/langchain_ai/models/ark_model.py index 8e48bf3d..d2f6ce0f 100644 --- a/veadk/community/langchain_ai/models/ark_model.py +++ b/veadk/community/langchain_ai/models/ark_model.py @@ -27,7 +27,7 @@ def __init__(self, model: str, **kwargs): model=model, api_key=secret_from_env("MODEL_AGENT_API_KEY")(), base_url=from_env( - "MODEL_AGENT_API_BASE_URL", + "MODEL_AGENT_API_BASE", default="https://ark.cn-beijing.volces.com/api/v3", )(), **kwargs,