|
|
import sys |
|
|
sys.path.append("../") |
|
|
from bw_utils import * |
|
|
from modules.embedding import get_embedding_model |
|
|
from langchain_experimental.generative_agents import GenerativeAgentMemory |
|
|
from langchain.retrievers import TimeWeightedVectorStoreRetriever |
|
|
from langchain_community.embeddings import HuggingFaceEmbeddings |
|
|
from langchain_community.llms import Tongyi,OpenAI |
|
|
from langchain_community.docstore import InMemoryDocstore |
|
|
from langchain_community.vectorstores import FAISS |
|
|
import faiss |
|
|
import math |
|
|
|
|
|
def build_role_agent_memory(type = "ga",**kwargs): |
|
|
if type == "ga": |
|
|
llm_name = kwargs["llm_name"] |
|
|
embedding_name = kwargs["embedding_name"] |
|
|
db_name = kwargs["db_name"] |
|
|
language = kwargs["language"] if "language" in kwargs else "" |
|
|
embedding_model = get_embedding_model(embedding_name, language=language) |
|
|
index = faiss.IndexFlatL2(len(embedding_model.embed_query("hello world"))) |
|
|
vectorstore = FAISS( |
|
|
embedding_function=embedding_model, |
|
|
index=index, |
|
|
docstore=InMemoryDocstore(), |
|
|
index_to_docstore_id={}, |
|
|
) |
|
|
memory_retriever = TimeWeightedVectorStoreRetriever(vectorstore=vectorstore, other_score_keys=["importance"], k=5) |
|
|
if llm_name.startswith("qwen"): |
|
|
chat_model = Tongyi( |
|
|
temperature=0.9, |
|
|
) |
|
|
else: |
|
|
chat_model = OpenAI( |
|
|
temperature=0.9, |
|
|
model="gpt-3.5-turbo", |
|
|
) |
|
|
agent_memory = RoleMemory_GA( |
|
|
llm=chat_model, |
|
|
memory_retriever=memory_retriever, |
|
|
embedding_model=embedding_model, |
|
|
memory_decay_rate=0.01 |
|
|
|
|
|
) |
|
|
return agent_memory |
|
|
|
|
|
else: |
|
|
db_name = kwargs["db_name"] |
|
|
embedding = kwargs["embedding"] |
|
|
db_type = kwargs["db_type"] if "db_type" in kwargs else "chromadb" |
|
|
capacity= kwargs["capacity"] if "capacity" in kwargs else 5 |
|
|
agent_memory = RoleMemory(db_name=db_name, |
|
|
embedding=embedding, |
|
|
db_type=db_type, |
|
|
capacity=capacity) |
|
|
return agent_memory |
|
|
|
|
|
def relevance_score_fn(score: float) -> float: |
|
|
return 1.0 - score / math.sqrt(2) |
|
|
|
|
|
class RoleMemory_GA(GenerativeAgentMemory): |
|
|
def init_from_data(self,data): |
|
|
for text in data: |
|
|
self.add_record(text) |
|
|
|
|
|
def add_record(self,text): |
|
|
self.add_memory(text) |
|
|
|
|
|
def search(self,query,top_k): |
|
|
fetched_memories = [doc.page_content for doc in self.fetch_memories(query)[:top_k]] |
|
|
if len(fetched_memories)>=top_k: |
|
|
print("-Memory Searching...") |
|
|
print(fetched_memories) |
|
|
return fetched_memories |
|
|
|
|
|
def delete_record(self, idx): |
|
|
pass |
|
|
|
|
|
|
|
|
class RoleMemory: |
|
|
def __init__(self,db_name,embedding,db_type = "chroma",capacity = 5,) -> None: |
|
|
self.idx = 0 |
|
|
self.capacity = capacity |
|
|
self.db_name = db_name |
|
|
self.db = build_db([],db_name,db_type,embedding,save_type="temporary") |
|
|
|
|
|
def init_from_data(self,data): |
|
|
for text in data: |
|
|
self.add_record(text) |
|
|
|
|
|
def add_record(self,text): |
|
|
self.idx += 1 |
|
|
self.db.add(text, str(self.idx), db_name=self.db_name) |
|
|
|
|
|
def search(self,query,top_k): |
|
|
return self.db.search(query, top_k,self.db_name) |
|
|
|
|
|
def delete_record(self, idx): |
|
|
self.db.delete(idx) |
|
|
|
|
|
@property |
|
|
def len(self): |
|
|
return self.db.len |
|
|
|
|
|
|
|
|
|
