Skip to content

Commit ab6e8f2

Browse files
committed
feat(arkitect): responses
1 parent 553d5b4 commit ab6e8f2

34 files changed

+21050
-1
lines changed
Lines changed: 28 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,28 @@
1+
# Copyright 2025 Bytedance 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 .base_memory_service import BaseMemoryService
16+
from .in_memory_memory_service import (
17+
InMemoryMemoryService,
18+
InMemoryMemoryServiceSingleton,
19+
)
20+
from .mem0_memory_service import Mem0MemoryService, Mem0MemoryServiceSingleton
21+
22+
__all__ = [
23+
"BaseMemoryService",
24+
"InMemoryMemoryService",
25+
"InMemoryMemoryServiceSingleton",
26+
"Mem0MemoryService",
27+
"Mem0MemoryServiceSingleton",
28+
]
Lines changed: 58 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,58 @@
1+
# Copyright 2025 Bytedance 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 abc import ABC, abstractmethod
16+
from typing import Any
17+
18+
from openai.types.responses import Response
19+
from pydantic import BaseModel
20+
from volcenginesdkarkruntime.types.chat.chat_completion_message import (
21+
ChatCompletionMessage,
22+
)
23+
24+
from arkitect.types.llm.model import ArkMessage
25+
26+
27+
class Memory(BaseModel):
28+
memory_content: str
29+
reference: Any | None = None
30+
metadata: Any | None = None
31+
32+
33+
class SearchMemoryResponse(BaseModel):
34+
memories: list[Memory]
35+
36+
@property
37+
def content(self) -> str:
38+
return "\n".join([m.memory_content for m in self.memories])
39+
40+
41+
class BaseMemoryService(ABC):
42+
@abstractmethod
43+
async def update_memory(
44+
self,
45+
user_id: str,
46+
new_messages: list[ArkMessage | dict | Response | ChatCompletionMessage],
47+
**kwargs: Any,
48+
) -> None:
49+
pass
50+
51+
@abstractmethod
52+
async def search_memory(
53+
self,
54+
user_id: str,
55+
query: str,
56+
**kwargs: Any,
57+
) -> SearchMemoryResponse:
58+
pass
Lines changed: 125 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,125 @@
1+
# Copyright 2025 Bytedance 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 typing import Any
16+
17+
from openai.types.responses import Response
18+
from typing_extensions import override
19+
from volcenginesdkarkruntime import AsyncArk
20+
from volcenginesdkarkruntime.types.chat.chat_completion_message import (
21+
ChatCompletionMessage,
22+
)
23+
24+
from arkitect.core.component.memory.base_memory_service import (
25+
BaseMemoryService,
26+
Memory,
27+
SearchMemoryResponse,
28+
)
29+
from arkitect.core.component.memory.utils import format_message_as_string
30+
from arkitect.types.llm.model import ArkMessage
31+
from arkitect.utils.common import Singleton
32+
33+
DEFAULT_SEARCH_MEM_PROMPT = """
34+
You have obtained a series of interactions between a user and an AI assistant.
35+
Please identify the user’s profile and other key information from
36+
past interactions to help answer the user’s new question.
37+
"""
38+
39+
DEFAULT_SEARCH_LLM_MODEL = "doubao-1-5-pro-32k-250115"
40+
41+
42+
class InMemoryMemoryService(BaseMemoryService):
43+
def __init__(
44+
self,
45+
default_search_model: str = DEFAULT_SEARCH_LLM_MODEL,
46+
default_search_prompt: str = DEFAULT_SEARCH_MEM_PROMPT,
47+
) -> None:
48+
self.default_search_model = default_search_model
49+
self.default_search_prompt = default_search_prompt
50+
51+
self.memory: dict = {}
52+
self._cached_query: dict = {}
53+
self._llm = AsyncArk()
54+
55+
@override
56+
async def update_memory(
57+
self,
58+
user_id: str,
59+
new_messages: list[ArkMessage | dict | Response | ChatCompletionMessage],
60+
**kwargs: Any,
61+
) -> None:
62+
if user_id not in self.memory:
63+
self.memory[user_id] = []
64+
self.memory[user_id].extend(new_messages)
65+
# invalidate cache
66+
self._cached_query[user_id] = {}
67+
68+
@override
69+
async def search_memory(
70+
self,
71+
user_id: str,
72+
query: str,
73+
**kwargs: Any,
74+
) -> SearchMemoryResponse:
75+
if user_id not in self.memory:
76+
return SearchMemoryResponse(
77+
memories=[
78+
Memory(
79+
memory_content="no memory found for this user",
80+
reference=None,
81+
)
82+
]
83+
)
84+
if self._cached_query.get(user_id, {}).get(query, None) is not None:
85+
return self._cached_query[user_id][query]
86+
memories = self.memory[user_id]
87+
results = "用户过去的交互记录\n\n"
88+
for memory in memories:
89+
content = format_message_as_string(memory)
90+
results += content
91+
summary = await self._llm.chat.completions.create(
92+
model=self.default_search_model,
93+
messages=[
94+
{
95+
"role": "system",
96+
"content": self.default_search_prompt,
97+
},
98+
{
99+
"role": "user",
100+
"content": results,
101+
},
102+
],
103+
stream=False,
104+
)
105+
memory_response = SearchMemoryResponse(
106+
memories=[
107+
Memory(
108+
memory_content=summary.choices[0].message.content,
109+
reference=None,
110+
)
111+
]
112+
)
113+
if user_id not in self._cached_query:
114+
self._cached_query[user_id] = {}
115+
self._cached_query[user_id][query] = memory_response
116+
return memory_response
117+
118+
@override
119+
async def delete_user(self, user_id: str) -> None:
120+
if user_id in self.memory:
121+
del self.memory[user_id]
122+
123+
124+
class InMemoryMemoryServiceSingleton(InMemoryMemoryService, Singleton):
125+
pass
Lines changed: 138 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,138 @@
1+
# Copyright 2025 Bytedance 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+
import asyncio
16+
import os
17+
from typing import Any
18+
19+
from mem0 import AsyncMemory as Mem0Memory
20+
from mem0.configs.base import MemoryConfig as Mem0Config
21+
from mem0.embeddings.configs import EmbedderConfig
22+
from mem0.llms.configs import LlmConfig
23+
from mem0.vector_stores.configs import VectorStoreConfig
24+
from openai.types.responses import Response
25+
from pydantic import BaseModel
26+
from typing_extensions import override
27+
from volcenginesdkarkruntime import AsyncArk
28+
from volcenginesdkarkruntime.types.chat.chat_completion_message import (
29+
ChatCompletionMessage,
30+
)
31+
32+
from arkitect.core.component.memory.base_memory_service import (
33+
BaseMemoryService,
34+
Memory,
35+
SearchMemoryResponse,
36+
)
37+
from arkitect.core.component.memory.utils import format_message_as_dict
38+
from arkitect.telemetry.logger import ERROR, INFO
39+
from arkitect.types.llm.model import ArkMessage
40+
from arkitect.utils.common import Singleton
41+
42+
DEFAULT_EMBEDDING_MODEL = "doubao-embedding-text-240715"
43+
DEFAULT_LLM_MODEL = "doubao-1-5-vision-pro-32k-250115"
44+
DEFAULT_BASE_URL = "https://ark.cn-beijing.volces.com/api/v3"
45+
46+
47+
default_ark_config = Mem0Config(
48+
embedder=EmbedderConfig(
49+
provider="openai",
50+
config={
51+
"model": DEFAULT_EMBEDDING_MODEL,
52+
"openai_base_url": DEFAULT_BASE_URL,
53+
"api_key": os.getenv("ARK_API_KEY"),
54+
"embedding_dims": 2560,
55+
},
56+
),
57+
llm=LlmConfig(
58+
provider="openai",
59+
config={
60+
"model": DEFAULT_LLM_MODEL,
61+
"openai_base_url": DEFAULT_BASE_URL,
62+
"api_key": os.getenv("ARK_API_KEY"),
63+
},
64+
),
65+
vector_store=VectorStoreConfig(config={"embedding_model_dims": 2560}),
66+
)
67+
68+
69+
class Mem0MemoryService(BaseMemoryService):
70+
71+
def __init__(self, config: Mem0Config = default_ark_config) -> None:
72+
self.mem0_config = config if config else Mem0Config()
73+
self._llm = AsyncArk()
74+
self.memory = Mem0Memory(config=self.mem0_config)
75+
self._task_queue: asyncio.Queue = asyncio.Queue()
76+
77+
@override
78+
async def update_memory(
79+
self,
80+
user_id: str,
81+
new_messages: list[ArkMessage | dict | Response | ChatCompletionMessage],
82+
blocking: bool = False,
83+
**kwargs: Any,
84+
) -> None:
85+
conversation = []
86+
for item in new_messages:
87+
conversation.append(format_message_as_dict(item))
88+
if blocking:
89+
await self._add_memory(conversation, user_id)
90+
else:
91+
await self._task_queue.put(
92+
asyncio.create_task(self._add_memory(conversation, user_id))
93+
)
94+
INFO("Memory update submitted")
95+
96+
async def _add_memory(self, conversation, user_id):
97+
await self.memory.add(conversation, user_id=user_id)
98+
INFO("Memory update completed")
99+
100+
async def _background_processor(self) -> None:
101+
while True:
102+
task = await self._task_queue.get()
103+
try:
104+
await task
105+
except Exception as e:
106+
ERROR(f"Memory update failed: {e}")
107+
self._task_queue.task_done()
108+
109+
@override
110+
async def search_memory(
111+
self,
112+
user_id: str,
113+
query: str,
114+
**kwargs: Any,
115+
) -> SearchMemoryResponse:
116+
relevant_memories = await self.memory.search(
117+
query=query, user_id=user_id, limit=3
118+
)
119+
fetched_results = relevant_memories.get("results", [])
120+
memeory_string = ""
121+
for element in fetched_results:
122+
memeory_string += element.get("memory", "") + "\n"
123+
return SearchMemoryResponse(
124+
memories=[
125+
Memory(
126+
memory_content=memeory_string,
127+
reference=None,
128+
metadata=relevant_memories,
129+
)
130+
]
131+
)
132+
133+
async def delete_user(self, user_id: str) -> None:
134+
await self.memory.delete_all(user_id=user_id)
135+
136+
137+
class Mem0MemoryServiceSingleton(Mem0MemoryService, Singleton):
138+
pass

0 commit comments

Comments
 (0)