[AgentShield] 진행 상황 및 개인 공부 정리 (2): 파이프라인, DB 연결, 데이터 축적

이번 글은 AgentShield 프로젝트를 진행하면서 정리한 개인 기록이다.

이번 단계에서 집중한 내용은 크게 세 가지다.

1. 공격 프롬프트를 실행하고 판정하는 파이프라인 구축
2. 실행 결과를 PostgreSQL과 ChromaDB에 저장하는 흐름 확인
3. 이후 Judge/Red Agent 파인튜닝을 위한 데이터 축적 구조 정리

AgentShield는 단순히 “LLM이 위험한 답변을 하는지” 확인하는 프로젝트가 아니라, 공격 생성 → 타겟 챗봇 실행 → 취약 여부 판정 → 결과 저장 → 재사용 및 파인튜닝으로 이어지는 LLM 보안 검증 파이프라인을 만드는 것이 목표다.

---

1. 현재까지 구현한 흐름

현재 파이프라인의 기본 흐름은 다음과 같다.

seed -> mutate -> target -> judge -> store -> reuse/fine-tune

각 단계의 의미는 다음과 같다.

• seed: 초기 공격 프롬프트 데이터
• mutate: Red Agent가 공격 프롬프트를 변형
• target: 타겟 LLM 또는 테스트베드 챗봇에 공격 실행
• judge: 응답을 보고 공격 성공 여부 판정
• store: 결과를 PostgreSQL, ChromaDB 등에 저장
• reuse/fine-tune: 축적된 데이터를 재사용하거나 파인튜닝 데이터로 활용

처음에는 규칙 기반 Judge로 시작했지만, 최종적으로는 Judge LLM과 Red Agent LLM을 학습시켜 더 정교한 보안 검증 구조로 확장하는 것이 목표다.

---

2. 파이프라인 실행

기본 파이프라인 실행 명령어는 다음과 같다.

cd /Users/parkyeonggon/Projects/final_project/AgentShield
source venv/bin/activate
python -m backend.graph.run_pipeline

특정 카테고리와 반복 횟수를 지정해서 실행할 수도 있다.

python -m backend.graph.run_pipeline -c LLM01 -m 5 -r 3
python -m backend.graph.run_pipeline -c LLM02 -m 5 -r 3

여기서 LLM01은 OWASP LLM Top 10 기준 Prompt Injection 항목을 의미한다.

---

3. ChromaDB 저장 확인

파이프라인이 끝난 뒤 공격 패턴이 ChromaDB에 저장되었는지 확인한다.

python -c "
from backend.rag.chromadb_client import get_rag_client
c = get_rag_client()
count = c.attack_col.count()
print(f'저장된 공격 패턴: {count}건')
if count > 0:
    docs = c.attack_col.peek(3)
    for d, m in zip(docs['documents'], docs['metadatas']):
        print(f'  [{m.get(\"category\",\"?\")}] {d[:80]}...')
"

ChromaDB는 단순 로그 저장소라기보다, 이후 Red Agent나 RAG 기반 탐색에서 재사용할 수 있는 공격 패턴 메모리 역할을 한다.

즉, PostgreSQL이 구조화된 실행 결과를 저장하는 DB라면, ChromaDB는 공격 문장과 패턴을 임베딩 기반으로 검색하고 재사용하기 위한 벡터 DB 역할을 한다.

---

4. PostgreSQL 연결 및 저장 확인

Docker Compose를 사용하는 경우 DB 컨테이너만 먼저 실행한다.

docker compose up -d db

DB 연결 상태를 확인한다.

docker compose exec db pg_isready -U agentshield -d agentshield

스키마가 적용되었는지 확인한다.

docker compose exec db psql -U agentshield -d agentshield -c "\dt"

파이프라인을 다시 실행하면 결과가 DB에 저장된다.

python -m backend.graph.run_pipeline -c LLM01 -m 1 -r 1

저장 결과는 다음 명령어로 확인할 수 있다.

/opt/homebrew/opt/postgresql@16/bin/psql -U agentshield -d agentshield -c "select count(*) from test_sessions;"

/opt/homebrew/opt/postgresql@16/bin/psql -U agentshield -d agentshield -c "select count(*) from test_results;"

/opt/homebrew/opt/postgresql@16/bin/psql -U agentshield -d agentshield -c "select session_id, phase, category, judgment, round, seed_id from test_results order by created_at desc limit 20;"

이 구조를 통해 단순 실행 로그가 아니라, 세션 단위의 공격 실험 결과를 계속 누적할 수 있다.

---

5. 로컬 PostgreSQL 설치 방식

Docker가 아니라 로컬 PostgreSQL을 사용할 경우 다음 순서로 설정한다.

brew install postgresql@16
brew services start postgresql@16

DB와 유저를 생성한다.

/opt/homebrew/opt/postgresql@16/bin/createuser -s agentshield 2>/dev/null
/opt/homebrew/opt/postgresql@16/bin/createdb -O agentshield agentshield 2>/dev/null

비밀번호를 설정한다.

/opt/homebrew/opt/postgresql@16/bin/psql -d agentshield -c "ALTER USER agentshield WITH PASSWORD 'agentshield';"

스키마를 적용한다.

/opt/homebrew/opt/postgresql@16/bin/psql -U agentshield -d agentshield -f database/schema.sql

연결 상태를 확인한다.

/opt/homebrew/opt/postgresql@16/bin/pg_isready -U agentshield -d agentshield

PostgreSQL 서비스 상태 확인과 중지는 다음 명령어로 할 수 있다.

brew services list | grep postgresql
brew services stop postgresql@16

---

6. 데이터 축적과 파인튜닝 방향

현재 목표는 공격 프롬프트, 타겟 응답, Judge 판정 결과를 계속 축적하는 것이다.

초기에는 R2가 만든 공격 프롬프트 데이터셋을 사용했다. 예를 들어 시스템 프롬프트 유출 요청, 고객 데이터 요청, 내부 정책 우회 요청 등 다양한 공격 시나리오를 포함한다.

현재 구조는 다음처럼 동작한다.

1. 공격 프롬프트를 타겟 LLM에 보낸다.
2. 타겟 LLM이 응답한다.
3. Judge가 응답을 분석해 공격 성공 여부를 판정한다.
4. 방어에 성공한 경우 Red Agent가 변형 공격을 생성한다.
5. 결과를 DB와 벡터 DB에 저장한다.
6. 저장된 데이터를 이후 학습 데이터로 재사용한다.

이후 목표는 규칙 기반 Judge를 넘어서는 것이다.

현재:

규칙 기반 Judge

목표:

AI Judge + Fine-tuned Judge + Red Agent 고도화

즉, 단순 패턴 매칭이 아니라 문맥과 의미를 보고 공격 성공 여부를 판단하는 Judge를 만들고, Red Agent도 더 정교한 변형 공격을 생성하도록 학습시키는 것이 목표다.

---

7. Ollama 모델 실험

로컬 환경에서는 Ollama 모델을 바꿔가며 Red Agent 또는 Target LLM으로 사용할 수 있다.

예시:

OLLAMA_RED_MODEL=gemma4:26b python -m backend.graph.run_pipeline

무검열 계열 모델을 실험할 수도 있다.

OLLAMA_RED_MODEL=huihui_ai/qwen2.5-abliterate:32b-instruct-q4_K_M python -m backend.graph.run_pipeline

모델 다운로드:

ollama pull huihui_ai/qwen2.5-abliterate:32b-instruct-q4_K_M

VRAM이 부족한 경우 더 작은 모델을 사용할 수 있다.

ollama pull huihui_ai/qwen2.5-abliterate:14b-instruct-q4_K_M

.env에서는 다음처럼 모델을 교체하며 실험할 수 있다.

OLLAMA_MODEL=gemma4:e2b
OLLAMA_RED_MODEL=gemma4:26b
OLLAMA_RED_MODEL=huihui_ai/qwen2.5-abliterate:32b-instruct-q4_K_M
OLLAMA_RED_MODEL=huihui_ai/qwen2.5-abliterate:14b-instruct-q4_K_M

여기서 중요한 점은 모델 자체의 성능뿐 아니라, 공격 생성 모델과 판정 모델의 역할을 분리해서 보는 것이다.

---

8. GPT 기반 Judge 실험 방향

초기에는 10개 정도의 샘플만 GPT로 판정해보는 방식이 적절하다고 판단했다.

이유는 다음과 같다.

• 전체 데이터를 GPT로 판정하면 비용이 커질 수 있다.
• 먼저 소량 데이터로 판정 품질을 확인해야 한다.
• 토큰 사용량과 응답 일관성을 확인해야 한다.
• 규칙 기반 Judge와 GPT Judge의 차이를 비교할 수 있다.

즉, 처음부터 모든 판정을 LLM에 맡기기보다, 규칙 기반 Judge와 LLM Judge를 비교하면서 점진적으로 확장하는 방식이 현실적이다.

---

9. 다음 작업

다음 작업은 다음과 같다.

1. 파이프라인을 반복 실행해 공격, 응답, 판정 데이터를 더 많이 축적한다.
2. ChromaDB에 저장된 공격 패턴을 Red Agent가 재사용할 수 있도록 연결한다.
3. PostgreSQL 기준으로 최신 세션 데이터를 export한다.
4. 규칙 기반 Judge와 GPT Judge의 판정 결과를 비교한다.
5. 잘못 판정된 케이스를 모아 Judge 개선 기준을 만든다.
6. Red Agent 변형 전략을 카테고리별로 분리한다.
7. 최종적으로 파인튜닝 가능한 데이터셋 형태로 정리한다.

데이터 export 예시는 다음과 같다.

python -m backend.finetuning.export_dpo_data --source db --session latest

JSON 결과 파일 기준으로 export할 수도 있다.

python -m backend.finetuning.export_dpo_data --source json

---

마무리

이번 단계에서 AgentShield는 단순한 LLM 보안 테스트 코드에서, 공격 실행과 판정 결과를 저장하고 재사용할 수 있는 파이프라인 구조로 넘어가고 있다.

현재 핵심은 “공격을 잘 만드는 것”만이 아니다.

공격을 실행하고, 응답을 판정하고, 결과를 저장하고, 다시 학습에 사용할 수 있는 구조를 만드는 것이다.

최종적으로는 다음 흐름을 목표로 한다.

Red Agent -> Target LLM -> Judge -> Store -> Fine-tune -> Re-test

 

이 구조가 완성되면 AgentShield는 LLM 보안 취약점을 반복적으로 검증하고 개선할 수 있는 테스트베드가 될 수 있다.