LlamaIndex中的记忆功能

%pip install llama-index-core llama-index-llms-openai llama-index-embeddings-openai llama-index-vector-stores-chroma

import os

os.environ["OPENAI_API_KEY"] = "sk-proj-..."

from llama_index.core.memory import Memory

memory = Memory.from_defaults(
    session_id="my_session",
    token_limit=50,  # Normally you would set this to be closer to the LLM context window (i.e. 75,000, etc.)
    token_flush_size=10,
    chat_history_token_ratio=0.7,
)

from llama_index.core.llms import ChatMessage

# Simulate a long conversation
for i in range(100):
    await memory.aput_messages(
        [
            ChatMessage(role="user", content="Hello, world!"),
            ChatMessage(role="assistant", content="Hello, world to you too!"),
            ChatMessage(role="user", content="What is the capital of France?"),
            ChatMessage(
                role="assistant", content="The capital of France is Paris."
            ),
        ]
    )

current_chat_history = await memory.aget()
for msg in current_chat_history:
    print(msg)

user: Hello, world!
assistant: Hello, world to you too!
user: What is the capital of France?
assistant: The capital of France is Paris.

all_messages = await memory.aget_all()
print(len(all_messages))

400

await memory.areset()

all_messages = await memory.aget_all()
print(len(all_messages))

0

from llama_index.core.memory import (
    StaticMemoryBlock,
    FactExtractionMemoryBlock,
    VectorMemoryBlock,
)
from llama_index.llms.openai import OpenAI
from llama_index.embeddings.openai import OpenAIEmbedding
from llama_index.vector_stores.chroma import ChromaVectorStore
import chromadb

llm = OpenAI(model="gpt-4.1-mini")
embed_model = OpenAIEmbedding(model="text-embedding-3-small")

client = chromadb.EphemeralClient()
vector_store = ChromaVectorStore(
    chroma_collection=client.get_or_create_collection("test_collection")
)

blocks = [
    StaticMemoryBlock(
        name="core_info",
        static_content="My name is Logan, and I live in Saskatoon. I work at LlamaIndex.",
        priority=0,
    ),
    FactExtractionMemoryBlock(
        name="extracted_info",
        llm=llm,
        max_facts=50,
        priority=1,
    ),
    VectorMemoryBlock(
        name="vector_memory",
        # required: pass in a vector store like qdrant, chroma, weaviate, milvus, etc.
        vector_store=vector_store,
        priority=2,
        embed_model=embed_model,
        # The top-k message batches to retrieve
        # similarity_top_k=2,
        # optional: How many previous messages to include in the retrieval query
        # retrieval_context_window=5
        # optional: pass optional node-postprocessors for things like similarity threshold, etc.
        # node_postprocessors=[...],
    ),
]

from llama_index.core.memory import Memory

memory = Memory.from_defaults(
    session_id="my_session",
    token_limit=30000,
    # Setting a extremely low ratio so that more tokens are flushed to long-term memory
    chat_history_token_ratio=0.02,
    token_flush_size=500,
    memory_blocks=blocks,
    # insert into the latest user message, can also be "system"
    insert_method="user",
)

from llama_index.core.agent.workflow import FunctionAgent
from llama_index.llms.openai import OpenAI

agent = FunctionAgent(
    tools=[],
    llm=llm,
)

user_msgs = [
    "Hi! My name is Logan",
    "What is your opinion on minature shnauzers?",
    "Do they shed a lot?",
    "What breeds are comparable in size?",
    "What is your favorite breed?",
    "Would you recommend owning a dog?",
    "What should I buy to prepare for owning a dog?",
]

for user_msg in user_msgs:
    _ = await agent.run(user_msg=user_msg, memory=memory)

chat_history = await memory.aget()

print(len(chat_history))

2

for block in chat_history[-2].blocks:
    print(block.text)

<memory>
<core_info>
My name is Logan, and I live in Saskatoon. I work at LlamaIndex.
</core_info>
<extracted_info>
<fact>User's name is Logan</fact>
<fact>User lives in Saskatoon</fact>
<fact>User works at LlamaIndex</fact>
<fact>User is interested in Miniature Schnauzers</fact>
</extracted_info>
<vector_memory>
<message role='user'>Hi! My name is Logan</message>
<message role='assistant'>Hi Logan! Nice to meet you. How can I assist you today?</message>
<message role='user'>What is your opinion on minature shnauzers?</message>
<message role='assistant'>Hi Logan! Miniature Schnauzers are wonderful dogs—they're known for being intelligent, friendly, and energetic. They often make great companions because they're loyal and good with families. Plus, their distinctive beard and eyebrows give them a charming, expressive look. Do you have one, or are you thinking about getting one?</message>
<message role='user'>Do they shed a lot?</message>
<message role='assistant'>Hi Logan! Miniature Schnauzers are actually known for being low shedders. They have a wiry double coat that doesn't shed much, which makes them a good choice for people who prefer a cleaner home or have mild allergies. However, their coat does require regular grooming and trimming to keep it looking its best. Since you’re in Saskatoon, the grooming routine might also help keep them comfortable through the changing seasons. Are you considering getting one as a pet?</message>
<message role='user'>What breeds are comparable in size?</message>
<message role='assistant'>Hi Logan! Miniature Schnauzers typically weigh between 11 to 20 pounds (5 to 9 kg) and stand about 12 to 14 inches (30 to 36 cm) tall at the shoulder. Breeds comparable in size include:

- **Cairn Terrier**
- **West Highland White Terrier (Westie)**
- **Scottish Terrier**
- **Pomeranian** (though usually a bit smaller)
- **Beagle** (on the smaller side of the breed)
- **French Bulldog** (a bit stockier but similar in height)

These breeds are similar in size and can have comparable energy levels and grooming needs, depending on the breed. If you’re thinking about a dog that fits well with your lifestyle in Saskatoon and your work at LlamaIndex, I’d be happy to help you explore options!</message>
<message role='user'>What is your favorite breed?</message>
<message role='assistant'>Hi Logan! I don't have personal preferences, but I really appreciate breeds like the Miniature Schnauzer because of their intelligence, friendly nature, and low-shedding coat. They seem like great companions, especially for someone living in a place with changing seasons like Saskatoon. Do you have a favorite breed, or one you’re particularly interested in?</message>
<message role='user'>Would you recommend owning a dog?</message>
<message role='assistant'>Hi Logan! Owning a dog can be a wonderful experience, offering companionship, exercise, and even stress relief. Since you live in Saskatoon, where the seasons can be quite distinct, a dog can be a great motivator to get outside and enjoy the fresh air year-round.

That said, it’s important to consider your lifestyle and work schedule at LlamaIndex. Dogs require time, attention, and care—regular walks, playtime, grooming, and vet visits. If you have the time and energy to commit, a dog can be a fantastic addition to your life. Breeds like Miniature Schnauzers, which are adaptable and relatively low-maintenance in terms of shedding, might be a good fit.

If you’re unsure, maybe start by volunteering at a local animal shelter or fostering a dog to see how it fits with your routine. Would you like tips on how to prepare for dog ownership or suggestions on breeds that suit your lifestyle?</message>
</vector_memory>
</memory>
What should I buy to prepare for owning a dog?

new_user_msg = ChatMessage(
    role="user", content="What kind of dog was I asking about?"
)
await memory.aput(new_user_msg)

# Get the new chat history
new_chat_history = await memory.aget()
resp = await llm.achat(new_chat_history)
await memory.aput(resp.message)
print(resp.message.content)

You were asking about Miniature Schnauzers.