mirror of
https://github.com/All-Hands-AI/OpenHands.git
synced 2026-01-09 14:57:59 -05:00
fb1822123a
* don't modify directories * oops typo * dev_config/python * add config to CI * bump CI python to 3.10 * 3.11? * del actions/ * add suggestions * delete unused code * missed some * oops missed another one * remove a file
39 lines
1.1 KiB
Python
39 lines
1.1 KiB
Python
from . import json
|
|
|
|
import chromadb
|
|
|
|
from llama_index.core import Document
|
|
from llama_index.core.retrievers import VectorIndexRetriever
|
|
from llama_index.core import VectorStoreIndex
|
|
from llama_index.vector_stores.chroma import ChromaVectorStore
|
|
|
|
class LongTermMemory:
|
|
def __init__(self):
|
|
db = chromadb.Client()
|
|
self.collection = db.create_collection(name="memories")
|
|
vector_store = ChromaVectorStore(chroma_collection=self.collection)
|
|
self.index = VectorStoreIndex.from_vector_store(vector_store)
|
|
self.thought_idx = 0
|
|
|
|
def add_event(self, event):
|
|
doc = Document(
|
|
text=json.dumps(event),
|
|
doc_id=self.thought_idx,
|
|
extra_info={
|
|
"type": event.action,
|
|
"idx": self.thought_idx,
|
|
},
|
|
)
|
|
self.thought_idx += 1
|
|
self.index.insert(doc)
|
|
|
|
def search(self, query, k=10):
|
|
retriever = VectorIndexRetriever(
|
|
index=self.index,
|
|
similarity_top_k=k,
|
|
)
|
|
results = retriever.retrieve(query)
|
|
return [r.get_text() for r in results]
|
|
|
|
|