Agentic AI/Memory

9. 장기기억(6) - 벡터스토어 만들고 테스트하기

아톨 2026. 9. 6. 20:45
반응형

이제 vector_store.py를 만들고, PostgreSQL의 memories 테이블과 ChromaDB를 연결하는 단계입니다.

 

[요약]

# 1. vector_store.py 만들기
# 2. 테스트 파일 만들기: test_vector_store.py
# 3. 테스트1.
# 4. 추가 수정
# 4.1.  검색 결과를 보기 좋게 출력하는 함수 추가: 
    ===> def print_search_results(results):
# 4.2.  database.py의 get_memories() 반환값 확인후 수정:
    ===> user_id 추가
# 4.3. PostgreSQL tuple을 ChromaDB용 dictionary로 변환하는 과정 추가:
    ===> def memory_tuple_to_dict(row):를 vector_store.py에 추가하기

 

다만 이번 단계에서는 agent_with_memory.py까지 한꺼번에 수정하지 않는 것이 좋습니다.

먼저 다음 구조를 완성하겠습니다.

PostgreSQL
    ↓
기억 저장
    ↓
ChromaDB
    ↓
의미 기반 검색
    ↓
검색된 기억 확인
 

이 단계가 성공하면 그다음에 기존 Agent의 search_memory()를 ChromaDB 기반 검색으로 교체하면 됩니다.


1. 이번 단계에서 만들 최종 구조

현재 파일은 다음과 같습니다.

agent_project/
│
├── .env
├── agent_with_memory.py
├── database.py
├── memory.py
│
├── test_chroma.py       ← ChromaDB 기본 테스트 완료
│
└── vector_store.py      ← 이번에 새로 작성
 

이번 단계가 끝나면 역할은 다음처럼 나뉩니다.

파일역할
database.py PostgreSQL 연결, 사용자 및 기억의 원본 저장
memory.py 기억 추출, 기억 저장 흐름
vector_store.py ChromaDB 저장 및 의미 검색
agent_with_memory.py 사용자 질문부터 답변까지 전체 실행
test_chroma.py ChromaDB 기본 기능 테스트

핵심은 다음입니다.

database.py는 데이터를 관리하고, vector_store.py는 데이터를 검색하기 쉽게 만듭니다.


2. 왜 PostgreSQL과 ChromaDB를 함께 사용해야 할까?

예를 들어 PostgreSQL에 다음 기억이 저장되어 있다고 하겠습니다.

memory_id: 15
user_id: 1
memory_type: goal
content: 사용자는 Python과 AI Agent를 학습하고 있다.
importance: 8
 

이 데이터를 ChromaDB에도 저장합니다.

ChromaDB id:
memory_15

document:
사용자는 Python과 AI Agent를 학습하고 있다.

metadata:
user_id = 1
memory_id = 15
memory_type = goal
importance = 8
 

그러면 두 데이터베이스는 다음처럼 연결됩니다.

PostgreSQL                         ChromaDB
─────────────                     ─────────────
memory_id = 15  ←──────────────→  id = "memory_15"
content          ←──────────────→ document
user_id          ←──────────────→ metadata.user_id
 

이 연결이 중요한 이유는 다음과 같습니다.

PostgreSQL은 원본

  • 기억 내용
  • 사용자 정보
  • 중요도
  • 생성일
  • 수정 및 삭제
  • 데이터 무결성

ChromaDB는 검색용 복사본

  • 문장의 임베딩
  • 의미 기반 검색
  • 유사도 검색
  • 사용자별 필터링

따라서 ChromaDB에서 검색한 결과를 PostgreSQL의 memory_id와 연결할 수 있습니다.


3. 이번 단계의 목표

이번에는 다음 세 가지 기능만 만들겠습니다.

 
get_collection()
 
 
add_memory()
 
 
search_similar_memories()
 

각각의 역할은 다음과 같습니다.

get_collection()
    ↓
ChromaDB 컬렉션 준비
 
add_memory()
    ↓
PostgreSQL의 기억을 ChromaDB에 저장
 
search_similar_memories()
    ↓
질문과 의미가 비슷한 기억 검색
 

4. vector_store.py 전체 코드

먼저 새 파일을 만드세요.

vector_store.py
 

그리고 다음 코드를 작성합니다.

 
    # ============================================================
    # vector_store.py
    # PostgreSQL의 기억을 ChromaDB에 저장하고 검색하는 모듈
    # ============================================================
    from pathlib import Path
    import chromadb

    # ============================================================
    # 1. ChromaDB 설정
    # ============================================================
    CHROMA_DIR = Path("./chroma_db")
    COLLECTION_NAME = "memories"

    # ============================================================
    # 2. ChromaDB 컬렉션 가져오기
    # ============================================================
    def get_collection():
        """ChromaDB의 memories 컬렉션을 가져온다. 컬렉션이 없으면 새로 생성한다."""
        client = chromadb.PersistentClient(path=str(CHROMA_DIR))
        collection = client.get_or_create_collection(name=COLLECTION_NAME)
        return collection

    # ============================================================
    # test_vector_store.py 만들고 수정(3)
    # ============================================================
    def memory_tuple_to_dict(row):
        """
        PostgreSQL에서 조회한 기억 tuple을 ChromaDB 저장용 dictionary로 변환한다.

        tuple 구조:
            (memory_id, user_id, memory_type, content, importance, created_at)
        """
        return {
            "memory_id": row[0],
            "user_id": row[1],
            "memory_type": row[2],
            "content": row[3],
            "importance": row[4],
            "created_at": row[5],
        }


    # ============================================================
    # 3. 기억 하나를 ChromaDB에 저장
    # ============================================================
    def add_memory(
        memory_id,
        user_id,
        content,
        memory_type=None,
        importance=5
    ):
        """
        기억 하나를 ChromaDB에 저장한다.
        Parameters
        ----------
        memory_id: int
            PostgreSQL memories 테이블의 memory_id
        user_id: int
            PostgreSQL user_profiles 테이블의 user_id
        content: str
            실제 기억 내용
        memory_type: str
            기억 유형
        importance: int
            기억 중요도
        """
        collection = get_collection()
        chroma_id = f"memory_{memory_id}"
        collection.upsert(
            ids=[chroma_id],
            documents=[content],
            metadatas=[
                {
                    "memory_id":memory_id,
                    "user_id":user_id,
                    "memory_type":memory_type or "general",
                    "importance":importance
                }
            ]
        )
        print(f"☑️ChromaDB 저장 완료: {chroma_id}")

    # ============================================================
    # 4. 여러 기억을 ChromaDB에 저장
    # ============================================================
    def add_memories(memories):
        """
        여러 기억을 한번에 ChromaDB에 저장한다.
        memories는 다음과 같은 dictionary 목록을 기대한다.
        [
            {
                "memory_id": 1,
                "user_id": 1,
                "content": "...",
                "memory_type": "goal",
                "importance": 8
            }
        ]
        """
        if not memories:
            return

        collection = get_collection()
        ids=[]
        documents=[]
        metadatas=[]

        for memory in memories:
            memory_id = memory["memory_id"]
            ids.append(f"memory_{memory_id}")
            documents.append(memory["content"])
            metadatas.append(
                {
                    "memory_id":memory_id,
                    "user_id":memory["user_id"],
                    "memory_type":memory.get("memory_type", "general"),
                    "importance":memory.get("importance", 5)
                }
            )
        collection.upsert(
            ids=ids,
            documents=documents,
            metadatas=metadatas
        )
        print(f"☑️ChromaDB에 {len(memories)}개 기억 저장 완료")

    # ============================================================
    # 5. 의미 기반 기억 검색
    # ============================================================
    def search_similar_memories(
            query,
            user_id=None,
            limit=5
    ):
        """
        query와 의미가 비슷한 기억을 검색한다.
        user_id가 지정되면 해당 사용자의 기억만 검색한다.    
        """
        collection = get_collection()
        query_kwargs = {
            "query_texts":[query],
            "n_results":limit
        }
        if user_id is not None:
            query_kwargs["where"] = {"user_id":user_id}
        results = collection.query(**query_kwargs)

        return results
 

5. 코드 하나씩 이해하기

5-1. ChromaDB 저장 경로

 
CHROMA_DIR = Path("./chroma_db")
 

이것은 ChromaDB 데이터를 저장할 폴더입니다.

agent_project/
│
├── chroma_db/
│   ├── ...
│
└── vector_store.py
 

PersistentClient를 사용하므로 프로그램이 종료되어도 데이터가 유지됩니다.

 
client = chromadb.PersistentClient(
    path=str(CHROMA_DIR)
)
 

여기서 str()을 사용하는 이유는 Path 객체를 문자열 경로로 변환하기 위해서입니다.


5-2. 컬렉션 가져오기

 
collection = client.get_or_create_collection(
    name=COLLECTION_NAME
)
 

이 코드는 다음을 의미합니다.

memories 컬렉션이 있으면
    → 기존 컬렉션 사용

없으면
    → 새 컬렉션 생성
 

따라서 프로그램을 실행할 때마다 새 컬렉션이 만들어지는 것은 아닙니다.


6. 왜 add_memory()에서 upsert()를 사용할까?

핵심 부분입니다.

 
chroma_id = f"memory_{memory_id}"
 

예를 들어 PostgreSQL의 memory_id가 15라면:

 
chroma_id = "memory_15"
 

그리고:

 
collection.upsert(
    ids=["memory_15"],
    documents=[
        "사용자는 Python과 AI Agent를 학습하고 있다."
    ],
    metadatas=[
        {
            "memory_id": 15,
            "user_id": 1,
            "memory_type": "goal",
            "importance": 8
        }
    ]
)
 

upsert()는 다음과 같습니다.

해당 ID가 없으면
    → 새로 저장

해당 ID가 있으면
    → 기존 데이터 갱신
 

따라서 같은 기억을 다시 저장해도 중복 오류가 발생하지 않습니다.

이것은 매우 중요합니다.

왜냐하면 나중에 Agent가 같은 기억을 다시 추출할 수 있기 때문입니다.


7. 왜 memory_id를 ChromaDB ID로 사용하는가?

다음과 같이 임의의 ID를 사용해도 됩니다.

 
ids=["abc123"]
 

하지만 우리는 PostgreSQL의 memory_id를 사용합니다.

 
ids=["memory_15"]
 

이렇게 하면 검색 결과를 PostgreSQL과 연결하기 쉽습니다.

ChromaDB 검색 결과
    ↓
memory_15
    ↓
memory_id = 15
    ↓
PostgreSQL memories 테이블 조회
 

즉:

ChromaDB는 검색
PostgreSQL은 원본 확인
 

이라는 구조가 됩니다.


8. 왜 metadata를 저장하는가?

다음 부분을 보세요.

 
metadatas=[
    {
        "memory_id": memory_id,
        "user_id": user_id,
        "memory_type": memory_type or "general",
        "importance": importance
    }
]
 

metadata는 문서에 붙이는 부가정보입니다.

예를 들어:

document:
사용자는 Python과 AI Agent를 학습하고 있다.

metadata:
user_id = 1
memory_id = 15
memory_type = goal
importance = 8
 

이렇게 하면 나중에 다음과 같은 검색이 가능합니다.

김현정 사용자의 기억만 검색
 
goal 유형의 기억만 검색
 
특정 사용자의 기억 중 중요도가 높은 기억 검색
 

이번 단계에서는 우선 user_id 필터만 사용하겠습니다.


9. search_similar_memories() 이해하기

핵심 코드는 다음입니다.

 
results = collection.query(
    query_texts=[
        query
    ],
    n_results=limit
)
 

예를 들어:

 
results = search_similar_memories(
    query="내가 요즘 어떤 공부를 하고 있지?",
    limit=3
)
 

ChromaDB는 질문을 임베딩으로 변환하고, 저장된 문서들과 의미적 유사도를 비교합니다.

질문:
내가 요즘 어떤 공부를 하고 있지?

검색 결과:
1. 사용자는 Python과 AI Agent를 학습하고 있다.
2. 사용자는 PostgreSQL 기반 장기기억을 구현했다.
3. 사용자는 ChromaDB를 학습하기 시작했다.
 

10. user_id 필터는 왜 필요한가?

다음 코드가 있습니다.

 
if user_id is not None:
    query_kwargs["where"] = {
        "user_id": user_id
    }
 

예를 들어:

 
search_similar_memories(
    query="공부",
    user_id=1,
    limit=5
)
 

그러면 ChromaDB는 다음 조건으로 검색합니다.

질문과 의미가 비슷하고
user_id가 1인 기억만 검색
 

이것은 매우 중요합니다.

사용자가 여러 명일 경우, 다른 사용자의 기억이 검색되면 안 되기 때문입니다.


11. 이제 테스트 파일을 만드세요

새 파일을 만듭니다.

test_vector_store.py
 

다음 코드를 작성합니다.

 
    # ============================================================
    # test_vector_store.py
    # vector_store.py 단독 테스트
    # ============================================================
    from vector_store import (
        add_memory,
        add_memories,
        search_similar_memories
    )

    def main():

        print("=" * 70)
        print("Vector Store 테스트")
        print("=" * 70)

        # --------------------------------------------------------
        # 1. 기억 하나 저장
        # --------------------------------------------------------  
        add_memory(
            memory_id=101,
            user_id=1,
            content="사용자는 Python과 AI Agent를 학습하고 있다.",
            memory_type="goal",
            importance=8
        )

        # --------------------------------------------------------
        # 2. 여러 기억 저장
        # --------------------------------------------------------
        memories = [
            {
                "memory_id":102,
                "user_id":1,
                "content":"사용자는 PostgreSQL을 Docker로 실행했다.",
                "memory_type":"project",
                "importance":7
            },
            {
                "memory_id":103,
                "user_id":1,
                "content":"사용자는 ChromaDB를 학습하기 시작했다.",
                "memory_type":"learning",
                "importance":8
            },
            {
                "memory_id":104,
                "user_id":2,
                "content":"사용자는 여행 계획을 세우고 있다.",
                "memory_type":"personal",
                "importance":5
            },
        ]

        add_memories(memories)

        # --------------------------------------------------------
        # 3. 전체 의미 검색
        # --------------------------------------------------------
        print("\n[전체 검색]")
        results = search_similar_memories(
            query="인공지능 에이전트 공부",
            limit=3
        )
        # print(results)
        print_search_results(results)

        # --------------------------------------------------------
        # 4. user_id별 의미 검색
        # --------------------------------------------------------
        print("\n[user_id=1 검색]")

        results = search_similar_memories(
            query="공부와 학습",
            user_id=1,
            limit=5
        )    
        # print(results)
        print_search_results(results)

        # --------------------------------------------------------
        # 5. 다른 사용자 검색
        # --------------------------------------------------------
        print("\n[user_id=2 검색]")

        results = search_similar_memories(
            query="공부와 학습",
            user_id=2,
            limit=5
        )  
        # print(results)
        print_search_results(results)

    # --------------------------------------------------------
    # 6. test_vector_store.py 만들고 수정(1)
    #  추가: 검색 결과를 보기 좋게 출력하는 함수:
    #  이경우 print(results)-> print_search_results(results)로 변경
    # --------------------------------------------------------
    def print_search_results(results):
        """  ChromaDB 검색 결과를 읽기 쉽게 출력한다."""
        ids = results.get("ids",[[]])[0]
        documents = results.get("documents",[[]])[0]    
        metadata = results.get("metadatas",[[]])[0]
        distances = results.get("distances",[[]])[0]

        for index in range(len(ids)):
            print("-"*60)
            print("ChromaDB ID: ", ids[index])
            print("문서: ", documents[index])
            print("메타데이터: ", metadata[index])
            print("거리: ", distances[index])

    if __name__=="__main__":
        main()


 

실행합니다.

 
python .\test_vector_store.py
 

12. 검색 결과는 어떻게 생기는가?

ChromaDB의 query() 결과는 일반적으로 다음과 같은 dictionary입니다.

 
{
    "ids": [
        [
            "memory_101",
            "memory_103",
            "memory_102"
        ]
    ],
    "distances": [
        [
            0.21,
            0.34,
            0.48
        ]
    ],
    "metadatas": [
        [
            {
                "memory_id": 101,
                "user_id": 1,
                "memory_type": "goal",
                "importance": 8
            },
            ...
        ]
    ],
    "documents": [
        [
            "사용자는 Python과 AI Agent를 학습하고 있다.",
            ...
        ]
    ]
}
 

여기서 중요한 점은 결과가 중첩 리스트라는 것입니다.

 
results["ids"][0]
 

첫 번째 질문에 대한 검색 결과 ID 목록입니다.

 
results["documents"][0]
 

첫 번째 질문에 대한 문서 목록입니다.

 
results["metadatas"][0]
 

첫 번째 질문에 대한 메타데이터 목록입니다.


13. 검색 결과를 보기 좋게 출력하는 함수

현재는 결과 전체를 출력하고 있습니다.

 
print(results)
 

하지만 나중에는 다음과 같이 출력하는 것이 더 좋습니다.

 
def print_search_results(results):
    """
    ChromaDB 검색 결과를 읽기 쉽게 출력한다.
    """

    ids = results.get("ids", [[]])[0]
    documents = results.get("documents", [[]])[0]
    metadatas = results.get("metadatas", [[]])[0]
    distances = results.get("distances", [[]])[0]

    for index in range(len(ids)):

        print("-" * 60)

        print("ChromaDB ID:", ids[index])

        print("문서:", documents[index])

        print("메타데이터:", metadatas[index])

        print("거리:", distances[index])
 

그리고 테스트 파일에서:

 
print_search_results(results)
 

를 사용하면 됩니다.


14. 거리(distance)는 어떻게 이해해야 할까?

ChromaDB 검색 결과에는 보통 distances가 포함됩니다.

예를 들어:

memory_101 → 0.21
memory_103 → 0.34
memory_102 → 0.48
 

일반적으로 같은 검색 결과 안에서는 거리가 작을수록 더 유사한 문서라고 이해하면 됩니다.

거리 0.21
    → 더 가까움
    → 더 유사할 가능성

거리 0.48
    → 더 멂
    → 상대적으로 덜 유사
 

다만 거리의 정확한 해석은 사용하는 임베딩 함수와 거리 계산 방식에 따라 달라질 수 있습니다.

이번 단계에서는 다음 정도만 기억하면 충분합니다.

distance는 검색어와 저장 문서가 얼마나 가까운지를 나타내는 값이다.


15. 중요한 수정 사항: memory.py의 반환 형식 확인

여기서 한 가지 중요한 부분이 있습니다.

현재 database.py의 get_memories()는 아마 PostgreSQL row를 tuple 형태로 반환할 것입니다.

예를 들어:

 
(
    15,
    1,
    "goal",
    "사용자는 Python과 AI Agent를 학습하고 있다.",
    8,
    datetime(...)
)
 

그런데 우리가 작성한 add_memories()는 dictionary를 기대합니다.

 
{
    "memory_id": 15,
    "user_id": 1,
    "content": "...",
    "memory_type": "goal",
    "importance": 8
}
 

따라서 실제 연결 단계에서는 PostgreSQL tuple을 ChromaDB용 dictionary로 변환하는 과정이 필요합니다.

이 부분은 다음처럼 처리할 수 있습니다.

 
def memory_tuple_to_dict(row):
    """
    PostgreSQL memories 테이블의 tuple을
    ChromaDB 저장용 dictionary로 변환한다.
    """

    return {
        "memory_id": row[0],
        "user_id": row[1],
        "memory_type": row[2],
        "content": row[3],
        "importance": row[4]
    }
 

단, 실제 get_memories()의 SELECT 순서가 다르면 인덱스도 달라집니다.

예를 들어 현재 SQL이 다음 순서라면:

 
SELECT
    memory_id,
    user_id,
    memory_type,
    content,
    importance,
    created_at
FROM memories
 

위 변환 함수가 맞습니다.


16. 이번 단계에서 아직 하지 않을 것

아직은 다음 작업을 하지 않는 것이 좋습니다.

agent_with_memory.py 전체 수정
 
memory.py의 search_memory() 교체
 
extract_memory()와 ChromaDB 연결
 
LLM 호출 구조 변경
 

이유는 간단합니다.

지금은 먼저 다음 질문에 답할 수 있어야 합니다.

“PostgreSQL에 있는 기억을 ChromaDB에 저장하고, 질문으로 관련 기억을 다시 찾을 수 있는가?”

이것이 독립적으로 확인된 후 Agent에 연결해야 합니다.


17. 이번 단계의 성공 기준

다음 네 가지가 모두 되면 성공입니다.

① PostgreSQL의 기억과 연결할 수 있는 ID가 있다

PostgreSQL memory_id = 101
ChromaDB id = memory_101
 

② ChromaDB에 기억이 저장된다

☑️ ChromaDB 저장 완료
 

③ 질문과 의미가 비슷한 기억이 검색된다

질문:
인공지능 에이전트 공부

검색 결과:
사용자는 Python과 AI Agent를 학습하고 있다.
 

④ 사용자별 검색이 된다

user_id=1
    → user_id가 1인 기억만 검색
 

18. 다음 단계의 전체 흐름

이번 vector_store.py 테스트가 성공하면 다음은 다음 구조로 발전합니다.

memory.py
    │
    ├── extract_memory()
    │       ↓
    │   중요한 기억 추출
    │
    ├── save_memory()
    │       ↓
    │   PostgreSQL 저장
    │       ↓
    │   memory_id 반환
    │
    └── add_memory()
            ↓
        ChromaDB 저장
 

검색은 다음처럼 바뀝니다.

사용자 질문
    ↓
ChromaDB 의미 검색
    ↓
관련 memory_id 확인
    ↓
PostgreSQL에서 상세 기억 조회
    ↓
build_memory_context()
    ↓
LLM에 전달
 

이것이 앞으로 만들 하이브리드 Memory Layer입니다.


지금 바로 진행할 순서

지금은 다음 순서로 하시면 됩니다.

① vector_store.py 작성
        ↓
② test_vector_store.py 작성
        ↓
③ 기억 저장 테스트
        ↓
④ 의미 검색 테스트
        ↓
⑤ user_id 필터 테스트
        ↓
⑥ PostgreSQL의 실제 memories 데이터와 연결
        ↓
⑦ memory.py에 ChromaDB 저장 기능 추가
        ↓
⑧ agent_with_memory.py의 검색 기능 교체
 

먼저 위 vector_store.py와 test_vector_store.py를 작성하고 실행해 보세요. 특히 test_vector_store.py에서 검색 결과가 정상적으로 나오는지 확인하는 것이 이번 단계의 핵심입니다.

다음에는 제가 현재 사용 중인 database.py, memory.py, agent_with_memory.py의 실제 코드를 기준으로 **PostgreSQL의 tuple을 ChromaDB에 연결하는 정확한 통합 코드**를 함께 완성해 드릴 수 있습니다.

반응형