import uuid
import numpy as np
from redis import Redis
from sentence_transformers import SentenceTransformer
import logging
from .gemini import generate_text
redis_client = Redis(host="localhost", port=6379)
# 임베딩 모델 초기화 (경량 모델 사용)
model = SentenceTransformer('all-MiniLM-L6-v2')
def save_document(title: str, url: str, content: str):
"""문서를 저장하고 생성된 vector_id를 반환"""
vector_id = str(uuid.uuid4())
embedding = model.encode(content)
vector = np.array(embedding, dtype=np.float32).tobytes()
redis_client.hset(
f"doc:{vector_id}",
mapping={
"title": title,
"url": url,
"content": content
}
)
return vector_id
def get_document(vector_id: str):
"""vector_id로 문서 조회"""
doc = redis_client.hgetall(f"doc:{vector_id}")
if not doc:
return None
return doc
def qa_with_redis(vector_id: str, query: str):
"""vector_id에 저장된 문서를 기반으로 답변 생성"""
doc = get_document(vector_id)
if not doc:
return "해당 문서를 찾을 수 없습니다."
content = doc[b"content"].decode() if isinstance(doc[b"content"], bytes) else doc[b"content"]
prompt = f"""
너는 공식 문서를 기반으로 질문에 답변하는 기술 문서 어시스턴트다.
규칙:
공식문서에 명시된 내용만 사용해서 답변한다.
공식문서에 없는 내용은 절대 추측하거나 보완 설명하지 않는다.
질문의 답이 공식문서에 없으면 "공식문서에 해당 내용이 없습니다."라고만 답변한다.
마크다운 문법, 줄바꿈 기호, 목록 표현, 기호 강조를 사용하지 않는다.
답변은 하나의 자연스러운 문장 또는 문단으로 작성한다.
공식문서:
{content}
사용자 질문:
{query}
"""
answer = generate_text(prompt)
logging.info(f"Vector ID {vector_id} - Generated answer: {answer}")
return answer