scriptoriumadmin/ReMe
1
0
Fork 0
mirror of https://github.com/agentscope-ai/ReMe.git synced 2026-09-06 08:16:00 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 1

update experience maker

Browse source
This commit is contained in:
jinli.yl 2025-07-08 11:25:18 +08:00
parent 3a2a2d1ae1
commit 20e02c0cf5
28 changed files with 940 additions and 277 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

2
experiencemaker/schema/experience.py
View file

@ -5,7 +5,7 @@ from uuid import uuid4
from pydantic import BaseModel, Field
from experiencemaker.schema.vector_store_node import VectorStoreNode
# https://docs.trychroma.com/docs/overview/getting-started
class ExperienceFunctionArg(BaseModel):
arg_name: str = Field(default=..., description="argument name")

1
experiencemaker/tool/base_tool.py
View file

@ -13,7 +13,6 @@ class BaseTool(BaseModel, ABC):
arguments: dict = Field(default_factory=dict, description="execute arguments")
enable_cache: bool = Field(default=False, description="whether to cache the tool result")
# TODO add cache expire
cached_result: dict = Field(default_factory=dict, description="tool execution result")
max_retries: int = Field(default=3, description="max retries")

13
v1/main.py → v1/app.py
View file

@ -4,10 +4,10 @@ from concurrent.futures.thread import ThreadPoolExecutor
import uvicorn
from fastapi import FastAPI
from v1.config.config_parser import ConfigParser
from v1.em_service import EMService
from v1.schema.request import RetrieverRequest, SummarizerRequest, VectorStoreRequest, AgentRequest
from v1.schema.response import RetrieverResponse, SummarizerResponse, VectorStoreResponse, AgentResponse
from v1.utils.config_parser import ConfigParser
app = FastAPI()
config_parser = ConfigParser(sys.argv[1:])
@ -17,28 +17,28 @@ thread_pool: ThreadPoolExecutor = ThreadPoolExecutor(max_workers=global_app_conf
@app.post('/retriever', response_model=RetrieverResponse)
def call_retriever(request: RetrieverRequest):
app_config = config_parser.get_app_config(**request.params)
app_config = config_parser.get_app_config(**request.config)
ems = EMService(app_config=app_config, thread_pool=thread_pool)
return ems.call_retriever(request)
@app.post('/summarizer', response_model=SummarizerResponse)
def call_summarizer(request: SummarizerRequest):
app_config = config_parser.get_app_config(**request.params)
app_config = config_parser.get_app_config(**request.config)
ems = EMService(app_config=app_config, thread_pool=thread_pool)
return ems.call_summarizer(request)
@app.post('/vector_store', response_model=VectorStoreResponse)
def call_vector_store(request: VectorStoreRequest):
app_config = config_parser.get_app_config(**request.params)
app_config = config_parser.get_app_config(**request.config)
ems = EMService(app_config=app_config, thread_pool=thread_pool)
return ems.call_vector_store(request)
@app.post('/agent', response_model=AgentResponse)
def call_agent(request: AgentRequest):
app_config = config_parser.get_app_config(**request.params)
app_config = config_parser.get_app_config(**request.config)
ems = EMService(app_config=app_config, thread_pool=thread_pool)
return ems.call_agent(request)
@ -48,4 +48,5 @@ if __name__ == "__main__":
host=global_app_config.http_service.host,
port=global_app_config.http_service.port,
timeout_keep_alive=global_app_config.http_service.timeout_keep_alive,
limit_concurrency=global_app_config.http_service.limit_concurrency)
limit_concurrency=global_app_config.http_service.limit_concurrency,
workers=global_app_config.http_service.workers)

28
v1/utils/config_parser.py → v1/config/config_parser.py
View file

@ -1,4 +1,5 @@
import json
from pathlib import Path
from loguru import logger
from omegaconf import OmegaConf, DictConfig
@ -6,24 +7,29 @@ from omegaconf import OmegaConf, DictConfig
from v1.schema.app_config import AppConfig
class ConfigParser(object):
class ConfigParser:
default_config_name: str = "demo_config.yaml"
def __init__(self, args: list):
assert len(args) >= 1, "The `args` require at least one argument."
assert "config_path=" in args[0], f"The 0th args must be config_path, for example: `config_path=XXX`."
config_path: str = args[0].replace("config_path=", "")
yaml_config = OmegaConf.load(config_path)
# step1: default config
self.app_config: DictConfig = OmegaConf.structured(AppConfig)
# step2: load from config yaml file
cli_config: DictConfig = OmegaConf.from_dotlist(args)
config_path = cli_config.get("config_path")
if config_path:
config_path = Path(config_path)
else:
config_path = Path(__file__).parent / self.default_config_name
logger.info(f"load config from path={config_path}")
yaml_config = OmegaConf.load(config_path)
self.app_config = OmegaConf.merge(self.app_config, yaml_config)
if args[1:]:
cli_config = OmegaConf.from_dotlist(args[1:])
self.app_config = OmegaConf.merge(self.app_config, cli_config)
# merge cli config
self.app_config = OmegaConf.merge(self.app_config, cli_config)
app_config_dict = OmegaConf.to_container(self.app_config, resolve=True)
app_config_str = json.dumps(app_config_dict, indent=2, ensure_ascii=False)
logger.info(f"app_config_str={app_config_str}")
logger.info(f"app_config_str={json.dumps(app_config_dict, indent=2, ensure_ascii=False)}")
def get_app_config(self, **kwargs) -> AppConfig:
app_config = self.app_config.copy()

15
v1/config.yaml → v1/config/demo_config.yaml
View file

@ -1,33 +1,20 @@
# demo config.yaml
#python -m experiencemaker.em_service \
# --port=8001 \
# --llm='{"backend": "openai_compatible", "model_name": "qwen3-32b", "temperature": 0.6}' \
# --embedding_model='{"backend": "openai_compatible", "model_name": "text-embedding-v4", "dimensions": 1024}' \
# --vector_store='{"backend": "elasticsearch"}' \
http_service:
host: "0.0.0.0"
port: 8001
timeout_keep_alive: 600
limit_concurrency: 64
# -http_service.port=8001
workers: 2
thread_pool:
max_workers: 20
api:
step_retriever: mock1_op->mock2_op->mock3_op
step_summarizer: mock1_op->[mock4_op->mock2_op|mock5_op]->mock3_op
vector_store: mock6_op
# -api.step_retriever=mock1_op->[mock4_op->mock2_op|mock5_op]->mock3_op
# -op.mock1_op.a=1
op:
mock1_op:
backend: mock1_op

4
v1/config_handler.py
View file

@ -1,4 +0,0 @@
class ConfigHandler(object):
def __init__(self):
pass

34
v1/em_service.py
View file

@ -1,5 +1,7 @@
from concurrent.futures import ThreadPoolExecutor
from loguru import logger
from v1.pipeline.pipeline import Pipeline
from v1.pipeline.pipeline_context import PipelineContext
from v1.schema.app_config import AppConfig
@ -33,30 +35,46 @@ class EMService(object):
self.context.set_context("request", request)
response = RetrieverResponse()
self.context.set_context("response", response)
pipeline = Pipeline(pipeline=self.context.app_config.api.retriever, context=self.context)
pipeline.execute_pipeline()
try:
Pipeline(pipeline=self.context.app_config.api.retriever, context=self.context)()
except Exception as e:
logger.exception(f"call_retriever encounter error={e.args}")
response.success = False
response.metadata["error"] = str(e)
return response
def call_summarizer(self, request: SummarizerRequest) -> SummarizerResponse:
self.context.set_context("request", request)
response = SummarizerResponse()
self.context.set_context("request", response)
pipeline = Pipeline(pipeline=self.context.app_config.api.summarizer, context=self.context)
pipeline.execute_pipeline()
try:
Pipeline(pipeline=self.context.app_config.api.summarizer, context=self.context)()
except Exception as e:
logger.exception(f"call_summarizer encounter error={e.args}")
response.success = False
response.metadata["error"] = str(e)
return response
def call_vector_store(self, request: VectorStoreRequest) -> VectorStoreResponse:
self.context.set_context("request", request)
response = VectorStoreResponse()
self.context.set_context("request", response)
pipeline = Pipeline(pipeline=self.context.app_config.api.vector_store, context=self.context)
pipeline.execute_pipeline()
try:
Pipeline(pipeline=self.context.app_config.api.vector_store, context=self.context)()
except Exception as e:
logger.exception(f"call_vector_store encounter error={e.args}")
response.success = False
response.metadata["error"] = str(e)
return response
def call_agent(self, request: AgentRequest) -> AgentResponse:
self.context.set_context("request", request)
response = AgentResponse()
self.context.set_context("request", response)
pipeline = Pipeline(pipeline=self.context.app_config.api.agent, context=self.context)
pipeline.execute_pipeline()
try:
Pipeline(pipeline=self.context.app_config.api.agent, context=self.context)()
except Exception as e:
logger.exception(f"call_agent encounter error={e.args}")
response.success = False
response.metadata["error"] = str(e)
return response

3
v1/embedding_model/__init__.py Normal file
View file

@ -0,0 +1,3 @@
from experiencemaker.utils.registry import Registry
EMBEDDING_MODEL_REGISTRY = Registry()

72
v1/embedding_model/base_embedding_model.py Normal file
View file

@ -0,0 +1,72 @@
from abc import ABC
from typing import List
from loguru import logger
from pydantic import BaseModel, Field
from v1.schema.vector_node import VectorNode
class BaseEmbeddingModel(BaseModel, ABC):
model_name: str = Field(default=..., description="model name")
dimensions: int = Field(default=..., description="dimensions")
max_retries: int = Field(default=3, description="max retries")
raise_exception: bool = Field(default=True, description="raise exception")
def _get_embeddings(self, input_text: str | List[str]):
"""
Get the embedding vector based on the input text.
This is an abstract method, and its concrete implementation must be provided in a subclass to generate the embedding vector for the given text.
Args:
input_text (str | List[str]): The input text, which can be a single string or a list of strings.
Raises:
NotImplementedError: If the method is not implemented in the subclass.
"""
raise NotImplementedError
def get_embeddings(self, input_text: str | List[str]):
"""
Retrieves embeddings for the input text.
This function attempts to obtain embeddings for the given input text. It will retry a maximum number of times in case of failure.
Parameters:
- input_text (str | List[str]): The input text, which can be a single string or a list of strings.
Returns:
- embeddings: The embeddings for the input text. Returns None if the maximum number of retries is reached and no successful result is obtained.
"""
# Attempt to get embeddings, with a maximum number of retries set
for i in range(self.max_retries):
try:
# Attempt to get embeddings, return immediately if successful
return self._get_embeddings(input_text)
except Exception as e:
# Log exception information when an error occurs
logger.exception(f"embedding model name={self.model_name} encounter error with e={e.args}")
# If the maximum number of retries is reached and raise_exception is set to True, re-throw the exception
if i == self.max_retries - 1 and self.raise_exception:
raise e
return None
def get_node_embeddings(self, nodes: VectorNode | List[VectorNode]):
if isinstance(nodes, VectorNode):
nodes.vector = self.get_embeddings(nodes.content)
return nodes
elif isinstance(nodes, list):
max_batch_size = 10 # text-embedding-v4 batch size should not be larger than 10
embeddings = [emb for i in range(0, len(nodes), max_batch_size) for emb in
self.get_embeddings(input_text=[node.content for node in nodes[i:i + max_batch_size]])]
if len(embeddings) != len(nodes):
logger.warning(f"embeddings.size={len(embeddings)} <> nodes.size={len(nodes)}")
else:
for node, embedding in zip(nodes, embeddings):
node.vector = embedding
return nodes
else:
raise RuntimeError(f"unsupported type={type(nodes)}")

92
v1/embedding_model/openai_compatible_embedding_model.py Normal file
View file

@ -0,0 +1,92 @@
import os
from typing import Literal, List
from openai import OpenAI
from pydantic import Field, PrivateAttr, model_validator
from v1.embedding_model import EMBEDDING_MODEL_REGISTRY
from v1.embedding_model.base_embedding_model import BaseEmbeddingModel
@EMBEDDING_MODEL_REGISTRY.register("openai_compatible")
class OpenAICompatibleEmbeddingModel(BaseEmbeddingModel):
api_key: str = Field(default_factory=lambda: os.getenv("OPENAI_API_KEY"), description="api key")
base_url: str = Field(default_factory=lambda: os.getenv("OPENAI_BASE_URL"), description="base url")
model_name: str = Field(default="text-embedding-v4", description="model name")
dimensions: int = Field(default=1024, description="dimensions")
encoding_format: Literal["float", "base64"] = Field(default="float", description="encoding_format")
_client: OpenAI = PrivateAttr()
@model_validator(mode="after")
def init_client(self):
"""
Initialize the OpenAI client after model validation.
This method is called after the model's data has been validated,
ensuring that all necessary attributes are correctly set before
initializing the OpenAI client. It creates an instance of the OpenAI
client using the provided API key and base URL, storing it in the
Returns:
self: Returns the instance of the current class for method chaining.
"""
self._client = OpenAI(api_key=self.api_key, base_url=self.base_url)
return self
def _get_embeddings(self, input_text: str | List[str]):
"""
Generate embeddings for the input text.
This method accepts either a single string or a list of strings as input,
and returns the corresponding embeddings based on the specified model.
Parameters:
- input_text (str | List[str]): The input text, which can be a single string or a list of strings.
Returns:
- List[float]: If the input is a single string, returns a list of floating-point numbers representing the embedding.
- List[List[float]]: If the input is a list of strings, returns a list where each item is a list of floating-point numbers representing the embedding of each string.
- Raises RuntimeError: If the input type is unsupported.
"""
# Create embeddings using the specified model, input text, dimensions, and encoding format
completion = self._client.embeddings.create(
model=self.model_name,
input=input_text,
dimensions=self.dimensions,
encoding_format=self.encoding_format
)
# Determine the type of input and process accordingly
if isinstance(input_text, str):
# If the input is a single string, return the embedding of that string
return completion.data[0].embedding
elif isinstance(input_text, list):
# If the input is a list of strings, initialize a list to hold the embeddings of each string
result_emb = [[] for _ in range(len(input_text))]
# Iterate through the generated embeddings and assign them to the corresponding positions in the result list
for emb in completion.data:
result_emb[emb.index] = emb.embedding
return result_emb
else:
# If the input type is neither a string nor a list of strings, throw an exception
raise RuntimeError(f"unsupported type={type(input_text)}")
def main():
from experiencemaker.utils.util_function import load_env_keys
load_env_keys()
model = OpenAICompatibleEmbeddingModel(dimensions=64, model_name="text-embedding-v4")
res1 = model.get_embeddings(
"The clothes are of good quality and look good, definitely worth the wait. I love them.")
res2 = model.get_embeddings(["aa", "bb"])
print(res1)
print(res2)
if __name__ == "__main__":
main()
# launch with: python -m experiencemaker.model.openai_compatible_embedding_model

70
v1/faiss_vector_store.py
View file

@ -1,70 +0,0 @@
import faiss
import numpy as np
from typing import List
from pydantic import Field, model_validator, PrivateAttr
from experiencemaker.model.base_embedding_model import BaseEmbeddingModel
from experiencemaker.schema.vector_store_node import VectorStoreNode
from experiencemaker.storage.base_vector_store import BaseVectorStore, VECTOR_STORE_REGISTRY
@VECTOR_STORE_REGISTRY.register("faiss")
class FaissVectorStore(BaseVectorStore):
dim: int = Field(default=768)
_index: faiss.IndexFlatIP = PrivateAttr()
_id_map: dict = PrivateAttr(default_factory=dict)
_vectors: list = PrivateAttr(default_factory=list)
@model_validator(mode="after")
def init_faiss(self):
self._index = faiss.IndexFlatIP(self.dim)
self._id_map = {}
self._vectors = []
return self
def exist_index(self, index_name: str | None = None) -> bool:
return hasattr(self, '_index')
def create_index(self, index_name: str | None = None):
self._index = faiss.IndexFlatIP(self.dim)
self._id_map = {}
self._vectors = []
def delete_index(self, index_name: str | None = None):
self._index = faiss.IndexFlatIP(self.dim)
self._id_map = {}
self._vectors = []
def exist_id(self, unique_id: str, index_name: str | None = None):
return unique_id in self._id_map
def insert(self, nodes: VectorStoreNode | List[VectorStoreNode], index_name: str | None = None, **kwargs):
if isinstance(nodes, VectorStoreNode):
nodes = [nodes]
nodes = self.embedding_model.get_node_embeddings(nodes)
for node in nodes:
vec = np.array(node.vector, dtype=np.float32)
self._index.add(vec.reshape(1, -1))
self._id_map[len(self._vectors)] = node.unique_id
self._vectors.append(node)
def update(self, nodes: VectorStoreNode | List[VectorStoreNode], index_name: str | None = None, **kwargs):
self.delete_index()
all_nodes = self._vectors + (nodes if isinstance(nodes, list) else [nodes])
self.insert(all_nodes)
def delete_by_id(self, unique_id: str, index_name: str | None = None, **kwargs):
idx_to_remove = [i for i, node in enumerate(self._vectors) if node.unique_id == unique_id]
if idx_to_remove:
self._vectors = [node for node in self._vectors if node.unique_id != unique_id]
self.create_index()
self.insert(self._vectors)
def retrieve_by_id(self, unique_id: str, index_name: str | None = None, **kwargs) -> VectorStoreNode | None:
for node in self._vectors:
if node.unique_id == unique_id:
return node
return None
def retrieve_by_query(self, query: str, top_k: int = 1, index_name: str | None = None, **kwargs) -> List[VectorStoreNode]:
query_vec = np.array(self.embedding_model.get_embeddings(query), dtype=np.float32).reshape(1, -1)
D, I = self._index.search(query_vec, top_k)
return [self._vectors[i] for i in I[0] if i >= 0]

3
v1/llm/__init__.py Normal file
View file

@ -0,0 +1,3 @@
from v1.utils.registry import Registry
LLM_REGISTRY = Registry()

52
v1/llm/base_llm.py Normal file
View file

@ -0,0 +1,52 @@
import time
from abc import ABC
from typing import List, Literal
from loguru import logger, Message
from pydantic import Field, BaseModel
from v1.tool.base_tool import BaseTool
class BaseLLM(BaseModel, ABC):
model_name: str = Field(...)
seed: int = Field(default=42)
top_p: float | None = Field(default=None)
# stream: bool = Field(default=True)
stream_options: dict = Field(default={"include_usage": True})
temperature: float = Field(default=0.0000001)
presence_penalty: float | None = Field(default=None)
enable_thinking: bool = Field(default=True, description="whether the current mode is the reasoning model, "
"or whether Qwen3's reasoning mode is currently enabled.")
tool_choice: Literal["none", "auto", "required"] = Field(default="auto", description="tool choice")
parallel_tool_calls: bool = Field(default=True)
max_retries: int = Field(default=5, description="max retries")
raise_exception: bool = Field(default=True, description="raise exception")
def stream_chat(self, messages: List[Message], tools: List[BaseTool] = None, **kwargs):
raise NotImplementedError
def stream_print(self, messages: List[Message], tools: List[BaseTool] = None, **kwargs):
raise NotImplementedError
def _chat(self, messages: List[Message], tools: List[BaseTool] = None, **kwargs) -> ActionMessage:
raise NotImplementedError
def chat(self, messages: List[Message], tools: List[BaseTool] = None, **kwargs) -> ActionMessage | None:
for i in range(self.max_retries):
try:
# Attempt to execute the chat logic
return self._chat(messages, tools, **kwargs)
except Exception as e:
# Log exceptions during the chat process
logger.exception(f"chat with model={self.model_name} encounter error with e={e.args}")
time.sleep(1 + i)
# If the maximum number of retries is reached and raise_exception is set to True, then re-throw the exception
if i == self.max_retries - 1 and self.raise_exception:
raise e
return None

7
v1/pipeline/pipeline.py
View file

@ -10,7 +10,7 @@ from v1.pipeline.pipeline_context import PipelineContext
from v1.utils.timer import Timer, timer
class Pipeline(object):
class Pipeline:
seq_symbol: str = "->"
parallel_symbol: str = "|"
@ -66,7 +66,8 @@ class Pipeline(object):
logger.info(f"stage_{i}: {op}")
elif isinstance(pipeline, list):
for op_list in zip_longest(*[self._parse_sub_pipeline(x) for x in pipeline], fillvalue="-"):
parallel_pipeline = [self._parse_sub_pipeline(x) for x in pipeline]
for op_list in zip_longest(*parallel_pipeline, fillvalue="-"):
i += 1
logger.info(f"stage{i}: {' | '.join(op_list)}")
@ -74,7 +75,7 @@ class Pipeline(object):
raise ValueError(f"unknown pipeline.type={type(pipeline)}")
@timer()
def execute_pipeline(self, enable_print: bool = True):
def __call__(self, enable_print: bool = True):
if enable_print:
self.print_pipeline()

8
v1/pipeline/pipeline_context.py
View file

@ -16,3 +16,11 @@ class PipelineContext(object):
def get_context(self, key: str):
return self.context.get(key)
@property
def request(self):
return self.get_context("request")
@property
def response(self):
return self.get_context("response")

2
v1/schema/app_config.py
View file

@ -8,6 +8,7 @@ class HttpServiceConfig:
port: int = field(default=8001)
timeout_keep_alive: int = field(default=600)
limit_concurrency: int = field(default=64)
workers: int | None = field(default=None)
@dataclass
@ -54,6 +55,7 @@ class VectorStoreConfig:
class AppConfig:
config_path: str = field(default="")
http_service: HttpServiceConfig = field(default_factory=HttpServiceConfig)
thread_pool: ThreadPoolConfig = field(default_factory=ThreadPoolConfig)
api: APIConfig = field(default_factory=APIConfig)

45
v1/schema/experience.py
View file

@ -5,7 +5,7 @@ from uuid import uuid4
from pydantic import BaseModel, Field
from v1.schema.vector_store_node import VectorStoreNode
from v1.schema.vector_node import VectorNode
class ExperienceMeta(BaseModel):
@ -14,26 +14,31 @@ class ExperienceMeta(BaseModel):
modified_time: str = Field(default_factory=lambda: datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S"))
extra_info: dict | None = Field(default=None)
def update_modified_time(self):
self.modified_time = datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")
class BaseExperienceNode(BaseModel, ABC):
workspace_id: str = Field(default="")
experience_id: str = Field(default_factory=lambda: uuid4().hex)
experience_type: str = Field(default="text")
workspace_id: str = Field(default="")
experience_meta: ExperienceMeta | None = Field(default=None)
def to_vector_store_node(self) -> VectorStoreNode:
...
@classmethod
def from_vector_store_node(cls, node: VectorStoreNode) -> "TextExperienceNode":
...
class TextExperienceNode(BaseExperienceNode):
experience_type: str = Field(default="text")
when_to_use: str = Field(default="")
content: str | bytes = Field(default="")
score: float | None = Field(default=None)
metadata: ExperienceMeta = Field(default_factory=ExperienceMeta)
def to_vector_node(self) -> VectorNode:
raise NotImplementedError
@classmethod
def from_vector_node(cls, node: VectorNode) -> "BaseExperienceNode":
raise NotImplementedError
class TextExperienceNode(BaseExperienceNode):
...
class FunctionArg(BaseModel):
@ -49,28 +54,16 @@ class Function(BaseModel):
class FuncExperienceNode(BaseExperienceNode):
"""
TODO
"""
experience_type: str = Field(default="func")
experience_type: str = Field(default="function")
functions: List[Function] = Field(default_factory=list)
class PersonalExperienceNode(BaseExperienceNode):
"""
TODO: memory node from MemoryScope
"""
experience_type: str = Field(default="personal")
person: str = Field(default="")
topic: str = Field(default="")
content: str | bytes = Field(default="")
class KnowledgeExperienceNode(BaseExperienceNode):
"""
TODO
"""
experience_type: str = Field(default="knowledge")
topic: str = Field(default="")
content: str | bytes = Field(default="")
score: float | None = Field(default=None)

4
v1/schema/request.py
View file

@ -7,8 +7,8 @@ from v1.schema.message import Message, Trajectory
class BaseRequest(BaseModel, ABC):
workspace_id: str = Field(default="")
params: dict = Field(default_factory=dict)
workspace_id: str = Field(default=...)
config: dict = Field(default_factory=dict)
metadata: dict | None = Field(default=None)

4
v1/schema/response.py
View file

@ -12,12 +12,12 @@ class BaseResponse(BaseModel):
class RetrieverResponse(BaseResponse):
experience_nodes: list[BaseExperienceNode] = Field(default_factory=list)
experience_nodes: List[BaseExperienceNode] = Field(default_factory=list)
experience_merged: str = Field(default="")
class SummarizerResponse(BaseResponse):
experience_nodes: list[BaseExperienceNode] = Field(default_factory=list)
experience_nodes: List[BaseExperienceNode] = Field(default_factory=list)
class VectorStoreResponse(BaseResponse):

2
v1/schema/vector_store_node.py → v1/schema/vector_node.py
View file

@ -3,7 +3,7 @@ from uuid import uuid4
from pydantic import BaseModel, Field
class VectorStoreNode(BaseModel):
class VectorNode(BaseModel):
unique_id: str = Field(default_factory=lambda: uuid4().hex)
workspace_id: str = Field(default="")
content: str = Field(default="")

113
v1/sqlite_vector_store.py
View file

@ -1,113 +0,0 @@
import sqlite3
import json
from typing import List
from pydantic import Field, model_validator, PrivateAttr
from experiencemaker.model.base_embedding_model import BaseEmbeddingModel
from experiencemaker.schema.vector_store_node import VectorStoreNode
from experiencemaker.storage.base_vector_store import BaseVectorStore, VECTOR_STORE_REGISTRY
@VECTOR_STORE_REGISTRY.register("sqlite")
class SQLiteVectorStore(BaseVectorStore):
db_path: str = Field(default="./vector_store.db")
_conn: sqlite3.Connection = PrivateAttr()
@model_validator(mode="after")
def init_db(self):
self._conn = sqlite3.connect(self.db_path, check_same_thread=False)
self._conn.execute(f"""
CREATE TABLE IF NOT EXISTS {self.index_name} (
unique_id TEXT PRIMARY KEY,
workspace_id TEXT,
content TEXT,
metadata TEXT,
vector TEXT
)
""")
self._conn.commit()
return self
def exist_index(self, index_name: str | None = None) -> bool:
index = index_name or self.index_name
cursor = self._conn.execute("SELECT name FROM sqlite_master WHERE type='table' AND name=?", (index,))
return cursor.fetchone() is not None
def create_index(self, index_name: str | None = None):
index = index_name or self.index_name
self._conn.execute(f"""
CREATE TABLE IF NOT EXISTS {index} (
unique_id TEXT PRIMARY KEY,
workspace_id TEXT,
content TEXT,
metadata TEXT,
vector TEXT
)
""")
self._conn.commit()
def delete_index(self, index_name: str | None = None):
index = index_name or self.index_name
self._conn.execute(f"DROP TABLE IF EXISTS {index}")
self._conn.commit()
def exist_id(self, unique_id: str, index_name: str | None = None):
index = index_name or self.index_name
cursor = self._conn.execute(f"SELECT 1 FROM {index} WHERE unique_id=?", (unique_id,))
return cursor.fetchone() is not None
def insert(self, nodes: VectorStoreNode | List[VectorStoreNode], index_name: str | None = None, **kwargs):
index = index_name or self.index_name
if isinstance(nodes, VectorStoreNode):
nodes = [nodes]
nodes = self.embedding_model.get_node_embeddings(nodes)
for node in nodes:
self._conn.execute(f"REPLACE INTO {index} (unique_id, workspace_id, content, metadata, vector) VALUES (?, ?, ?, ?, ?)",
(node.unique_id, node.workspace_id, node.content, json.dumps(node.metadata), json.dumps(node.vector)))
self._conn.commit()
def update(self, nodes: VectorStoreNode | List[VectorStoreNode], index_name: str | None = None, **kwargs):
self.insert(nodes, index_name=index_name)
def delete_by_id(self, unique_id: str, index_name: str | None = None, **kwargs):
index = index_name or self.index_name
self._conn.execute(f"DELETE FROM {index} WHERE unique_id=?", (unique_id,))
self._conn.commit()
def retrieve_by_id(self, unique_id: str, index_name: str | None = None, **kwargs) -> VectorStoreNode | None:
index = index_name or self.index_name
cursor = self._conn.execute(f"SELECT unique_id, workspace_id, content, metadata, vector FROM {index} WHERE unique_id=?", (unique_id,))
row = cursor.fetchone()
if row:
return VectorStoreNode(
unique_id=row[0],
workspace_id=row[1],
content=row[2],
metadata=json.loads(row[3]),
vector=json.loads(row[4])
)
return None
def retrieve_by_query(self, query: str, top_k: int = 1, index_name: str | None = None, **kwargs) -> List[VectorStoreNode]:
index = index_name or self.index_name
query_vec = self.embedding_model.get_embeddings(query)
cursor = self._conn.execute(f"SELECT unique_id, workspace_id, content, metadata, vector FROM {index}")
results = []
for row in cursor:
node = VectorStoreNode(
unique_id=row[0],
workspace_id=row[1],
content=row[2],
metadata=json.loads(row[3]),
vector=json.loads(row[4])
)
node.metadata["score"] = self._cosine_similarity(query_vec, node.vector)
results.append(node)
results.sort(key=lambda x: x.metadata["score"], reverse=True)
return results[:top_k]
@staticmethod
def _cosine_similarity(vec1, vec2):
import math
dot = sum(x * y for x, y in zip(vec1, vec2))
norm1 = math.sqrt(sum(x * x for x in vec1))
norm2 = math.sqrt(sum(y * y for y in vec2))
return dot / (norm1 * norm2) if norm1 and norm2 else 0.0

3
v1/tool/__init__.py Normal file
View file

@ -0,0 +1,3 @@
from v1.utils.registry import Registry
TOOL_REGISTRY = Registry()

79
v1/tool/base_tool.py Normal file
View file

@ -0,0 +1,79 @@
from abc import ABC
from loguru import logger
from pydantic import BaseModel, Field
class BaseTool(BaseModel, ABC):
tool_id: str = Field(default="")
name: str = Field(..., description="tool name")
description: str = Field(..., description="tool description")
tool_type: str = Field(default="function")
parameters: dict = Field(default_factory=dict, description="tool parameters")
arguments: dict = Field(default_factory=dict, description="execute arguments")
enable_cache: bool = Field(default=False, description="whether to cache the tool result")
cached_result: dict = Field(default_factory=dict, description="tool execution result")
max_retries: int = Field(default=3, description="max retries")
raise_exception: bool = Field(default=True, description="raise exception")
success: bool = Field(default=True, description="whether the tool executed successfully")
def reset(self):
self.arguments.clear()
self.success = True
def _execute(self, **kwargs):
raise NotImplementedError
def execute(self, **kwargs):
cache_id = ""
if self.enable_cache:
cache_id = self.get_cache_id(**kwargs)
if cache_id in self.cached_result:
return self.cached_result[cache_id]
for i in range(self.max_retries):
try:
if self.enable_cache:
self.cached_result[cache_id] = self._execute(**kwargs)
return self.cached_result[cache_id]
else:
return self._execute(**kwargs)
except Exception as e:
logger.exception(f"using tool.name={self.name} encounter error with e={e.args}")
if i == self.max_retries - 1 and self.raise_exception:
raise e
return None
def simple_dump(self) -> dict:
"""
It may be in other different tool params formats; different versions are completed here.
"""
return {
"type": self.tool_type,
self.tool_type: {
"name": self.name,
"description": self.description,
"parameters": self.parameters,
},
}
@property
def input_schema(self) -> dict:
return self.parameters.get("properties", {})
@property
def output_schema(self) -> dict:
raise NotImplementedError
def refresh(self):
# for mcp
raise NotImplementedError
def get_cache_id(self, **kwargs) -> str:
raise NotImplementedError

15
v1/tt.py
View file

@ -1,15 +0,0 @@
class Pipe():
def __init__(self):
pass
if __name__ == "__main__":
line = "a->[c->d->d|e->f]->[f->g]->h"
ops = []
for sub_line in line.split("["):
for sub_line2 in sub_line.split("]"):
ops.append(sub_line2)
print(ops)

3
v1/vector_store/__init__.py Normal file
View file

@ -0,0 +1,3 @@
from v1.utils.registry import Registry
VECTOR_STORE_REGISTRY = Registry()

44
v1/vector_store/base_vector_store.py Normal file
View file

@ -0,0 +1,44 @@
from abc import ABC
from typing import List
from pydantic import BaseModel, Field
from v1.model.base_embedding_model import BaseEmbeddingModel
from v1.schema.vector_node import VectorNode
class BaseVectorStore(BaseModel, ABC):
embedding_model: BaseEmbeddingModel = Field(default=...)
def exist_workspace(self, workspace_id: str, **kwargs) -> bool:
raise NotImplementedError
def delete_workspace(self, workspace_id: str, **kwargs):
raise NotImplementedError
def create_workspace(self, workspace_id: str, **kwargs):
raise NotImplementedError
def dump_workspace(self, workspace_id: str, **kwargs):
raise NotImplementedError
def load_workspace(self, workspace_id: str, **kwargs):
raise NotImplementedError
def retrieve_by_query(self, query: str, workspace_id: str, top_k: int = 1, **kwargs) -> List[VectorNode]:
raise NotImplementedError
def retrieve_by_id(self, unique_id: str, workspace_id: str = None, **kwargs) -> VectorNode | None:
raise NotImplementedError
def insert(self, nodes: VectorNode | List[VectorNode], workspace_id: str, **kwargs):
raise NotImplementedError
def update(self, nodes: VectorNode | List[VectorNode], workspace_id: str, **kwargs):
raise NotImplementedError
def exist_id(self, unique_id: str, workspace_id: str = None, **kwargs) -> bool:
raise NotImplementedError
def delete_id(self, unique_id: str, workspace_id: str = None, **kwargs):
raise NotImplementedError

290
v1/vector_store/es_vector_store.py Normal file
View file

@ -0,0 +1,290 @@
import os
from typing import List, Tuple
from elasticsearch import Elasticsearch
from elasticsearch.helpers import bulk
from loguru import logger
from pydantic import Field, PrivateAttr, model_validator
from v1.model.openai_compatible_embedding_model import OpenAICompatibleEmbeddingModel
from v1.schema.vector_store_node import VectorStoreNode
from v1.storage.base_vector_store import BaseVectorStore, VECTOR_STORE_REGISTRY
@VECTOR_STORE_REGISTRY.register("elasticsearch")
class EsVectorStore(BaseVectorStore):
hosts: str | List[str] = Field(default_factory=lambda: os.getenv("ES_HOSTS", "http://localhost:9200"))
basic_auth: str | Tuple[str, str] | None = Field(default=None)
bulk_chunk_size: int = Field(default=512)
retrieve_filters: List[dict] = []
_client: Elasticsearch = PrivateAttr()
@model_validator(mode="after")
def init_client(self):
if isinstance(self.hosts, str):
hosts = [self.hosts]
else:
hosts = self.hosts
self._client = Elasticsearch(hosts=hosts, basic_auth=self.basic_auth)
return self
def exist_index(self, index_name: str = None) -> bool:
if index_name is None:
index_name = self.index_name
return self._client.indices.exists(index=index_name)
def delete_index(self, index_name: str = None):
if index_name is None:
index_name = self.index_name
if self._client.indices.exists(index=index_name):
self._client.indices.delete(index=index_name)
def create_index(self, index_name: str = None):
if index_name is None:
index_name = self.index_name
if self._client.indices.exists(index=index_name):
logger.warning(f"index_name={index_name} is already exists!")
return None
index = {
"mappings": {
"properties": {
"workspace_id": {"type": "keyword"},
"content": {"type": "text"},
"metadata": {"type": "object"},
"vector": {
"type": "dense_vector",
"dims": self.embedding_model.dimensions
}
}
}
}
return self._client.indices.create(index=index_name, body=index)
def refresh_index(self, index_name: str = None):
if index_name is None:
index_name = self.index_name
self._client.indices.refresh(index=index_name)
@staticmethod
def doc2node(doc) -> VectorStoreNode:
node = VectorStoreNode(**doc["_source"])
node.unique_id = doc["_id"]
if "_score" in doc:
node.metadata["_score"] = doc["_score"] - 1
return node
def exist_id(self, unique_id: str, index_name: str = None):
if index_name is None:
index_name = self.index_name
return self._client.exists(index=index_name, id=unique_id)
def node2doc(self, node: VectorStoreNode, add_op_type: bool = False, index_name: str = None) -> dict:
if index_name is None:
index_name = self.index_name
doc: dict = {
"_index": index_name,
"_id": node.unique_id,
"_source": {
"workspace_id": node.workspace_id,
"content": node.content,
"metadata": node.metadata,
"vector": node.vector
}
}
if add_op_type:
doc["_op_type"] = "update" if self.exist_id(node.unique_id, index_name) else "index",
return doc
def add_term_filter(self, key: str, value):
if key:
self.retrieve_filters.append({"term": {key: value}})
return self
def add_range_filter(self, key: str, gte=None, lte=None):
if key:
if gte is not None and lte is not None:
self.retrieve_filters.append({"range": {key: {"gte": gte, "lte": lte}}})
elif gte is not None:
self.retrieve_filters.append({"range": {key: {"gte": gte}}})
elif lte is not None:
self.retrieve_filters.append({"range": {key: {"lte": lte}}})
return self
def clear_filter(self):
self.retrieve_filters.clear()
return self
def insert(self, nodes: VectorStoreNode | List[VectorStoreNode], refresh_index: bool = True, index_name: str = None,
**kwargs):
if index_name is None:
index_name = self.index_name
if not self.exist_index(index_name):
self.create_index(index_name)
if isinstance(nodes, VectorStoreNode):
nodes = [nodes]
embedded_nodes = [node for node in nodes if node.vector]
not_embedded_nodes = [node for node in nodes if not node.vector]
now_embedded_nodes = self.embedding_model.get_node_embeddings(not_embedded_nodes)
docs = [self.node2doc(node, False, index_name) for node in embedded_nodes + now_embedded_nodes]
status, error = bulk(self._client, docs, chunk_size=self.bulk_chunk_size, **kwargs)
logger.info(f"insert sample.size={len(nodes)} status={status} error={error}")
if refresh_index:
self.refresh_index(index_name)
def update(self, nodes: VectorStoreNode | List[VectorStoreNode], refresh_index: bool = True, index_name: str = None,
**kwargs):
if index_name is None:
index_name = self.index_name
if not self.exist_index(index_name):
self.create_index(index_name)
if isinstance(nodes, VectorStoreNode):
nodes = [nodes]
nodes = self.embedding_model.get_node_embeddings(nodes)
docs = [self.node2doc(node, True, index_name) for node in nodes]
status, error = bulk(self._client, docs, chunk_size=self.bulk_chunk_size, **kwargs)
update_size = sum([1 if doc["_op_type"] == "update" else 0 for doc in docs])
insert_size = len(docs) - update_size
logger.info(f"update update_size={update_size} insert_size={insert_size} status={status} error={error}")
if refresh_index:
self.refresh_index(index_name)
def delete_by_id(self, unique_id: str, index_name: str = None, **kwargs):
if index_name is None:
index_name = self.index_name
if not self.exist_index(index_name):
self.create_index(index_name)
return self._client.delete(index=index_name, id=unique_id, **kwargs)
def retrieve_by_id(self, unique_id: str, index_name: str = None, **kwargs) -> VectorStoreNode | None:
if index_name is None:
index_name = self.index_name
if not self.exist_index(index_name):
logger.warning(f"index_name={index_name} is not exists!")
return None
try:
doc = self._client.get(index=index_name, id=unique_id, **kwargs)
return self.doc2node(doc)
except Exception as e:
logger.warning(f"{index_name} retrieve_by_id unique_id={unique_id} is not found with error={e.args}")
return None
def retrieve_by_query(self, query: str, top_k: int = 1, index_name: str = None, **kwargs) -> List[VectorStoreNode]:
if index_name is None:
index_name = self.index_name
if not self.exist_index(index_name):
logger.warning(f"index_name={index_name} is not exists!")
return []
query_vector = self.embedding_model.get_embeddings(query)
body = {
"query": {
"script_score": {
"query": {"bool": {"must": self.retrieve_filters}},
"script": {
"source": "cosineSimilarity(params.query_vector, 'vector') + 1.0",
"params": {"query_vector": query_vector},
}
}
},
"size": top_k
}
response = self._client.search(index=index_name, body=body, **kwargs)
nodes: List[VectorStoreNode] = []
for doc in response['hits']['hits']:
nodes.append(self.doc2node(doc))
self.retrieve_filters.clear()
return nodes
def main():
from experiencemaker.utils.util_function import load_env_keys
load_env_keys()
embedding_model = OpenAICompatibleEmbeddingModel(dimensions=1024)
index_name = "rag_nodes_index"
hosts = "http://11.160.132.46:8200"
es = EsVectorStore(hosts=hosts, embedding_model=embedding_model, index_name=index_name)
es.delete_index()
es.create_index()
sample_nodes = [
VectorStoreNode(
workspace_id="w1",
content="Artificial intelligence is a technology that simulates human intelligence.",
metadata={
"node_type": "n1",
}
),
VectorStoreNode(
workspace_id="w1",
content="AI is the future of mankind.",
metadata={
"node_type": "n1",
}
),
VectorStoreNode(
workspace_id="w1",
content="I want to eat fish!",
metadata={
"node_type": "n2",
}
),
VectorStoreNode(
workspace_id="w2",
content="The bigger the storm, the more expensive the fish.",
metadata={
"node_type": "n1",
}
),
]
es.insert(sample_nodes, refresh_index=True)
logger.info("=" * 20)
results = es.add_term_filter(key="workspace_id", value="w1") \
.add_term_filter(key="metadata.node_type", value="n1") \
.retrieve_by_query("What is AI?", top_k=5)
for r in results:
logger.info(r.model_dump(exclude={"vector"}))
logger.info("=" * 20)
logger.info("=" * 20)
results = es.add_term_filter(key="workspace_id", value="w1") \
.retrieve_by_query("What is AI?", top_k=5)
for r in results:
logger.info(r.model_dump(exclude={"vector"}))
logger.info("=" * 20)
logger.info("=" * 20)
results = es.retrieve_by_query("What is AI?", top_k=5)
for r in results:
logger.info(r.model_dump(exclude={"vector"}))
logger.info("=" * 20)
if __name__ == "__main__":
main()
# launch with: python -m experiencemaker.storage.es_vector_store

209
v1/vector_store/file_vector_store.py Normal file
View file

@ -0,0 +1,209 @@
import json
import math
import threading
from pathlib import Path
from typing import List, Any
from loguru import logger
from pydantic import Field, model_validator, PrivateAttr
from experiencemaker.model.openai_compatible_embedding_model import OpenAICompatibleEmbeddingModel
from experiencemaker.schema.vector_store_node import VectorStoreNode
from experiencemaker.storage.base_vector_store import BaseVectorStore, VECTOR_STORE_REGISTRY
@VECTOR_STORE_REGISTRY.register("local_file")
class FileVectorStore(BaseVectorStore):
store_dir: str = Field(default="./file_vector_store")
index_path: Path | None = Field(default=None)
_thread_lock: Any = PrivateAttr()
@model_validator(mode="after")
def init_client(self):
self._thread_lock = threading.Lock()
store_path = Path(self.store_dir)
store_path.mkdir(parents=True, exist_ok=True)
self.index_path = store_path / f"{self.index_name}.jsonl"
if not self.index_path.exists():
self.index_path.touch(exist_ok=True)
return self
def get_index_path(self, index_name: str = None) -> Path:
if index_name is None:
index_path = self.index_path
else:
store_path = Path(self.store_dir)
index_path = store_path / f"{self.index_name}.jsonl"
if not index_path.exists():
index_path.touch(exist_ok=True)
return index_path
def exist_index(self, index_name: str = None) -> bool:
index_path = self.get_index_path(index_name)
with self._thread_lock:
return index_path.exists()
def delete_index(self, index_name: str = None):
index_path = self.get_index_path(index_name)
with self._thread_lock:
if index_path.exists() and index_path.is_file():
index_path.unlink()
def create_index(self, index_name: str = None):
index_path = self.get_index_path(index_name)
with self._thread_lock:
if not index_path.exists():
index_path.touch(exist_ok=True)
def _load(self, index_name: str = None) -> List[VectorStoreNode]:
index_path = self.get_index_path(index_name)
nodes = []
with self._thread_lock:
with open(index_path) as f:
for line in f:
if line.strip():
nodes.append(VectorStoreNode(**json.loads(line)))
return nodes
def _dump(self, nodes: List[VectorStoreNode], index_name: str = None):
index_path = self.get_index_path(index_name)
with self._thread_lock:
with open(index_path, "w") as f:
for doc in nodes:
f.write(doc.model_dump_json() + "\n")
def exist_id(self, unique_id: str, index_name: str = None):
nodes = self._load(index_name=index_name)
for node in nodes:
if node.unique_id == unique_id:
return True
return False
def insert(self, nodes: VectorStoreNode | List[VectorStoreNode], index_name: str = None, **kwargs):
if index_name is None:
index_name = self.index_name
return self.update(nodes, index_name=index_name, **kwargs)
def update(self, nodes: VectorStoreNode | List[VectorStoreNode], index_name: str = None, **kwargs):
if index_name is None:
index_name = self.index_name
if isinstance(nodes, VectorStoreNode):
nodes = [nodes]
all_node_dict = {}
nodes: List[VectorStoreNode] = self.embedding_model.get_node_embeddings(nodes)
exist_nodes: List[VectorStoreNode] = self._load(index_name=index_name)
for node in exist_nodes:
all_node_dict[node.unique_id] = node
update_cnt = 0
for node in nodes:
if node.unique_id in all_node_dict:
update_cnt += 1
all_node_dict[node.unique_id] = node
self._dump(list(all_node_dict.values()), index_name=index_name)
logger.info(
f"update {index_name} nodes.size={len(nodes)} all.size={len(all_node_dict)} update_cnt={update_cnt}")
def delete_by_id(self, unique_id: str, index_name: str = None, **kwargs):
if index_name is None:
index_name = self.index_name
nodes = self._load(index_name=index_name)
dump_nodes: List[VectorStoreNode] = []
for node in nodes:
if node.unique_id != unique_id:
dump_nodes.append(node)
if len(dump_nodes) < len(nodes):
self._dump(dump_nodes, index_name=index_name)
logger.info(f"delete_by_id unique_id={unique_id}")
def retrieve_by_id(self, unique_id: str, index_name: str = None, **kwargs) -> VectorStoreNode | None:
nodes = self._load(index_name=index_name)
for node in nodes:
if node.unique_id == unique_id:
return node
return None
@staticmethod
def calculate_similarity(query_vector: List[float], node_vector: List[float]):
assert query_vector, f"query_vector is empty!"
assert node_vector, f"node_vector is empty!"
assert len(query_vector) == len(node_vector), \
f"query_vector.size={len(query_vector)} node_vector.size={len(node_vector)}"
dot_product = sum(x * y for x, y in zip(query_vector, node_vector))
norm_v1 = math.sqrt(sum(x ** 2 for x in query_vector))
norm_v2 = math.sqrt(sum(y ** 2 for y in node_vector))
return dot_product / (norm_v1 * norm_v2)
def retrieve_by_query(self, query: str, top_k: int = 1, index_name: str = None, **kwargs) -> List[VectorStoreNode]:
query_vector = self.embedding_model.get_embeddings(query)
nodes: List[VectorStoreNode] = self._load(index_name=index_name)
for node in nodes:
node.metadata["score"] = self.calculate_similarity(query_vector, node.vector)
nodes = sorted(nodes, key=lambda x: x.metadata["score"], reverse=True)
return nodes[:top_k]
def main():
from experiencemaker.utils.util_function import load_env_keys
load_env_keys()
embedding_model = OpenAICompatibleEmbeddingModel(dimensions=1024)
index_name = "rag_nodes_index"
client = FileVectorStore(embedding_model=embedding_model, index_name=index_name)
client.delete_index()
client.create_index()
sample_nodes = [
VectorStoreNode(
workspace_id="w1",
content="Artificial intelligence is a technology that simulates human intelligence.",
metadata={
"node_type": "n1",
}
),
VectorStoreNode(
workspace_id="w1",
content="AI is the future of mankind.",
metadata={
"node_type": "n1",
}
),
VectorStoreNode(
workspace_id="w1",
content="I want to eat fish!",
metadata={
"node_type": "n2",
}
),
VectorStoreNode(
workspace_id="w2",
content="The bigger the storm, the more expensive the fish.",
metadata={
"node_type": "n1",
}
),
]
client.insert(sample_nodes)
logger.info("=" * 20)
results = client.retrieve_by_query("What is AI?", top_k=5)
for r in results:
logger.info(r.model_dump(exclude={"vector"}))
logger.info("=" * 20)
if __name__ == "__main__":
main()
# launch with: python -m experiencemaker.storage.file_vector_store