FAISS Memory

[amd-srijaroy]

First create "def get_embedding(text: str) -> list:" which will return a vector(array) from a string. Then use the below code

import os
import uuid
from typing import Any, Dict, List, Optional

import faiss
import numpy as np
from pydantic import BaseModel, Field


from autogen_core.memory import Memory, MemoryContent, MemoryMimeType

class FAISSVectorMemoryConfig(BaseModel):
    """
    Configuration parameters for FAISS-based vector memory.
    """
    dimension: int  = Field(..., description="Embedding dimension.")
    use_gpu: bool   = Field(False, description="Use GPU if available.")
    top_k: int      = Field(2, description="How many results to return on query.")
    score_threshold: float = Field(0.0, description="Minimum similarity score.")

class FAISSVectorMemory(Memory):
    """
    A simple example that uses an in-memory FAISS index to store embeddings and retrieve them.
    """

    def __init__(self, config: FAISSVectorMemoryConfig, *args, **kwargs) -> None:
        super().__init__(*args, **kwargs)
        self.config = config
        self.dimension = self.config.dimension

        # Create a FAISS index (flat L2, but customizable)
        self.index = faiss.IndexFlatL2(self.dimension)
        self.use_gpu = self.config.use_gpu

        self._embeddings: List[np.ndarray] = []
        self._documents: List[Dict[str, Any]] = []

        if self.use_gpu:
            # Move to GPU if supported
            res = faiss.StandardGpuResources()
            self.index = faiss.index_cpu_to_gpu(res, 0, self.index)

    async def add(self, content: MemoryContent) -> None:
        """
        Add a single document/chunk to the FAISS index.
        """
        # Prepare embedding using the function defined elsewhere
        embedding = get_embedding(content.content)
        embedding_np = np.array(embedding).astype("float32").reshape(1, -1)

        # Add to FAISS index
        self.index.add(embedding_np)

        # Store in local memory
        self._embeddings.append(embedding_np)
        self._documents.append(
            {
                "id": str(uuid.uuid4()),
                "content": content.content,
                "mime_type": content.mime_type,
                "metadata": content.metadata or {},
            }
        )

    async def query(self, input_query: MemoryContent) -> List[MemoryContent]:
        """
        Query FAISS for the most similar documents.
        """
        if not self._documents:
            return []

        # Compute query embedding using the function defined elsewhere
        query_embedding = get_embedding(input_query.content)
        query_embedding_np = np.array(query_embedding).astype("float32").reshape(1, -1)

        # Search in FAISS index
        distances, indices = self.index.search(query_embedding_np, self.config.top_k)
        
        results = []
        for dist, idx in zip(distances[0], indices[0]):
            if idx == -1:
                continue

            # Convert L2 distance to similarity score
            score = 1.0 / (1.0 + dist)
            if score < self.config.score_threshold:
                continue

            # Reconstruct matched document
            doc = self._documents[idx]
            results.append(
                MemoryContent(
                    content=doc["content"],
                    mime_type=doc["mime_type"],
                    metadata={**doc["metadata"], "score": float(score)},
                )
            )

        return results

    async def update_context(self, input_query: Any) -> None:
        """
        Update context based on matched documents.
        Extract `content` appropriately based on the context structure.
        """
        # Suppose context has an attribute `text` or similar; adjust as needed
        query_text = input_query.get('text') if hasattr(input_query, 'get') else str(input_query)

        # Generate embedding from the extracted text content
        query_embedding = get_embedding(query_text)
        query_embedding_np = np.array(query_embedding).astype("float32").reshape(1, -1)

        # Search in FAISS index
        distances, indices = self.index.search(query_embedding_np, self.config.top_k)

        results = []
        for dist, idx in zip(distances[0], indices[0]):
            if idx == -1:
                continue

            score = 1.0 / (1.0 + dist)
            if score < self.config.score_threshold:
                continue

            doc = self._documents[idx]
            results.append(
                MemoryContent(
                    content=doc["content"],
                    mime_type=doc["mime_type"],
                    metadata={**doc["metadata"], "score": float(score)},
                )
            )

        # Output or handle matched contents as needed in your agent's context
        for mc in results:
            print(f"Matched chunk with score {mc.metadata.get('score')}:\n{mc.content}\n")

    def dump_component(self) -> BaseModel:
        """
        Serialize the FAISS index to disk.
        """
        return self.config

    async def clear(self) -> None:
        """
        Clear all stored memory data.
        """
        self.index.reset()
        self._embeddings.clear()
        self._documents.clear()

    async def close(self):
        """
        Cleanup method if needed for persistence.
        """
        pass

Test

from autogen_agentchat.agents import AssistantAgent
from autogen_agentchat.ui import Console
from autogen_core.memory import MemoryContent, MemoryMimeType
# You're importing get_weather tool for agent usage
# from my_tools import get_weather  
from autogen_ext.models.openai import OpenAIChatCompletionClient

# Initialize FAISS memory
faiss_user_memory = FAISSVectorMemory(
    config=FAISSVectorMemoryConfig(
        dimension=1536,  # Replace with the correct dimension for your embeddings
        use_gpu=False,
        top_k=2,
        score_threshold=0.3,
    )
)

# Add user preferences to memory
await faiss_user_memory.add(
    MemoryContent(
        content="The weather should be in metric units, not imperial",
        mime_type=MemoryMimeType.TEXT,
        metadata={"category": "preferences", "type": "units"},
    )
)

await faiss_user_memory.add(
    MemoryContent(
        content="Meal recipe must be vegan",
        mime_type=MemoryMimeType.TEXT,
        metadata={"category": "preferences", "type": "dietary"},
    )
)


async def get_weather(city: str, units: str ) -> str:
    """

The code you provided defines an asynchronous function called `get_weather` that takes two parameters: `city` (a string representing the city name) and `units` (a string representing the unit system for temperature, either "imperial" or "metric"). The function returns a string that describes the weather in the specified city based on the unit system provided.
    
    
    """
    if units == "imperial":
        return f"The weather in {city} is 73 °F and Sunny."
    elif units == "metric":
        return f"The weather in {city} is 23 °C and Sunny."
    else:
        return f"Sorry, I don't know the weather in {city}."


# Create AssistantAgent with FAISS memory
assistant_agent = AssistantAgent(
    name="assistant_agent",
    model_client=model_client,  --> Create a LLM MODEL Client
    tools=[get_weather],  # Replace with your tools like get_weather
    memory=[faiss_user_memory],
)

# Run conversation
output = await assistant_agent.run(task="How is the weather in NewYork?")
print(output)

[junmiron]

Hey this is great!! by any chance do you have an example of using ChromaDBVectorMemory with embeddings, using an embedding_function in chromadb. Thanks!

[amd-srijaroy]

This might help

https://docs.trychroma.com/docs/embeddings/embedding-functions

[ekzhu]

@amd-srijaroy thanks would you like to create an extension package for your memory module and submit a PR to update the community extension list?

[amd-srijaroy]

Not sure. Please feel free to do if you want.