Spaces:
Paused
Paused
File size: 1,611 Bytes
e5902c3 deab47c e5902c3 deab47c e5902c3 15a2751 201c475 bab921d deab47c 15a2751 201c475 15a2751 e5902c3 201c475 15a2751 deab47c 15a2751 e5902c3 15a2751 e5902c3 15a2751 e5902c3 15a2751 e5902c3 2d2ee7a 15a2751 e5902c3 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 |
from agent import Agent
from python.helpers.vector_db import VectorDB, Document
from python.helpers import files
import os, json
from python.helpers.tool import Tool, Response
from python.helpers.print_style import PrintStyle
db: VectorDB | None = None
class Memory(Tool):
def execute(self,**kwargs):
result = process_query(self.agent, self.args["memory"],self.args["action"], result_count=self.agent.config.auto_memory_count)
return Response(message="\n\n".join(result), break_loop=False)
def initialize(embeddings_model, subdir=""):
global db
dir = os.path.join("memory",subdir)
db = VectorDB(embeddings_model=embeddings_model, in_memory=False, cache_dir=dir)
def process_query(agent:Agent, message: str, action: str = "load", result_count: int = 3, **kwargs):
if not db: initialize(agent.config.embeddings_model, subdir=agent.config.memory_subdir)
if action.strip().lower() == "save":
id = db.insert_document(str(message)) # type: ignore
return files.read_file("./prompts/fw.memory_saved.md")
elif action.strip().lower() == "delete":
deleted = db.delete_documents(message) # type: ignore
return files.read_file("./prompts/fw.memories_deleted.md", count=deleted)
else:
results=[]
docs = db.search_max_rel(message,result_count) # type: ignore
if len(docs)==0: return files.read_file("./prompts/fw.memories_not_found.md", query=message)
for doc in docs:
results.append(doc.page_content)
return results
# return "\n\n".join(results)
|