Skip to content

Commit 8724858

Browse files
committed
feat: support langchain and langgraph community
1 parent 549d486 commit 8724858

File tree

13 files changed

+372
-3
lines changed

13 files changed

+372
-3
lines changed
Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,13 @@
1+
# Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates.
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# http://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,13 @@
1+
# Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates.
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# http://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,13 @@
1+
# Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates.
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# http://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
Lines changed: 60 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,60 @@
1+
# Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates.
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# http://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
15+
from __future__ import annotations
16+
17+
from langchain.agents import AgentState
18+
from langchain.agents.middleware import after_agent
19+
from langchain_core.messages.ai import AIMessage
20+
from langchain_core.messages.human import HumanMessage
21+
from langgraph.runtime import Runtime
22+
23+
from veadk.community.langchain_ai.store.memory.viking_memory import (
24+
VikingMemoryStore,
25+
)
26+
from veadk.utils.logger import get_logger
27+
28+
logger = get_logger(__name__)
29+
30+
31+
@after_agent
32+
def save_session(state: AgentState, runtime: Runtime) -> None:
33+
"""Save the session to the memory store."""
34+
store: VikingMemoryStore | None = runtime.store
35+
if not store:
36+
return
37+
38+
app_name = store.index
39+
user_id = runtime.context.user_id
40+
session_id = runtime.context.session_id
41+
42+
messages = state.get("messages", [])
43+
logger.debug(
44+
f"Save session {session_id} for user {user_id} with {len(messages)} messages. messages={messages}"
45+
)
46+
47+
events = {}
48+
for message in messages:
49+
print(type(message))
50+
if isinstance(message, HumanMessage):
51+
event = {"role": "user", "parts": [{"text": message.content}]}
52+
53+
elif isinstance(message, AIMessage):
54+
event = {"role": "assistant", "parts": [{"text": message.content}]}
55+
else:
56+
...
57+
58+
events[message.id] = event
59+
60+
store.put(namespace=(app_name, user_id), key=session_id, value=events)
Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,13 @@
1+
# Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates.
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# http://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
Lines changed: 34 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,34 @@
1+
# Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates.
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# http://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
15+
from __future__ import annotations
16+
17+
from langchain_core.utils.utils import (
18+
from_env,
19+
secret_from_env,
20+
)
21+
from langchain_openai import ChatOpenAI
22+
23+
24+
class ArkChatModel(ChatOpenAI):
25+
def __init__(self, model: str, **kwargs):
26+
super().__init__(
27+
model=model,
28+
api_key=secret_from_env("MODEL_AGENT_API_KEY")(),
29+
base_url=from_env(
30+
"MODEL_AGENT_API_BASE_URL",
31+
default="https://ark.cn-beijing.volces.com/api/v3",
32+
)(),
33+
**kwargs,
34+
)
Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,13 @@
1+
# Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates.
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# http://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,13 @@
1+
# Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates.
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# http://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
Lines changed: 102 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,102 @@
1+
# Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates.
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# http://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
15+
from __future__ import annotations
16+
17+
import json
18+
from collections.abc import Iterable
19+
20+
from langgraph.store.base import (
21+
BaseStore,
22+
GetOp,
23+
ListNamespacesOp,
24+
Op,
25+
PutOp,
26+
Result,
27+
SearchOp,
28+
)
29+
30+
from veadk.memory.long_term_memory_backends.vikingdb_memory_backend import (
31+
VikingDBLTMBackend,
32+
)
33+
34+
35+
# mem0
36+
# viking db
37+
class VikingMemoryStore(BaseStore):
38+
def __init__(self, index: str):
39+
# index = (index, user_id)
40+
# key = session_id
41+
self.index = index
42+
self._backend = VikingDBLTMBackend(index=index)
43+
44+
def batch(self, ops: Iterable[Op]) -> list[Result]:
45+
# The batch/abatch methods are treated as internal.
46+
# Users should access via put/search/get/list_namespaces/etc.
47+
results = []
48+
for op in ops:
49+
if isinstance(op, PutOp):
50+
self._apply_put_op(op)
51+
elif isinstance(op, GetOp):
52+
self._apply_get_op(op)
53+
elif isinstance(op, SearchOp):
54+
results.extend(self._apply_search_op(op))
55+
# elif isinstance(op, ListNamespacesOp):
56+
# self._apply_list_namespaces_op(op)
57+
else:
58+
raise ValueError(f"Unknown op type: {type(op)}")
59+
60+
return results
61+
62+
def abatch(
63+
self, ops: Iterable[GetOp | SearchOp | PutOp | ListNamespacesOp]
64+
) -> list[Result]: ...
65+
66+
def _apply_put_op(self, op: PutOp) -> None:
67+
index, user_id = op.namespace
68+
session_id = op.key
69+
70+
assert index == self._backend.index, (
71+
"index must be the same as the backend index"
72+
)
73+
74+
value = op.value
75+
76+
event_strings = []
77+
78+
for _, event in value.items():
79+
event_strings.append(json.dumps(event))
80+
81+
if self._backend.save_memory(
82+
user_id=user_id,
83+
session_id=session_id,
84+
event_strings=event_strings,
85+
):
86+
return None
87+
88+
def _apply_get_op(self, op: GetOp):
89+
return ["Not implemented"]
90+
91+
def _apply_search_op(self, op: SearchOp):
92+
index, user_id = op.namespace_prefix
93+
assert index == self._backend.index, (
94+
"index must be the same as the backend index"
95+
)
96+
97+
query = op.query
98+
if not query:
99+
return []
100+
101+
value = self._backend.search_memory(user_id=user_id, query=query, top_k=1)
102+
return value
Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,13 @@
1+
# Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates.
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# http://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.

0 commit comments

Comments
 (0)