캡스톤 프로젝트: 엔터프라이즈 지식 에이전트

ROUTING_PROMPT = """
이 쿼리를 다음 중 하나로 분류하세요: FAQ, TECHNICAL, POLICY
 
쿼리: {query}
 
분류:
"""
 
def route_query(query: str) -> str:
    response = client.chat.completions.create(
        model="gpt-4o-mini",
        messages=[
            {"role": "system", "content": ROUTING_PROMPT.format(query=query)}
        ]
    )
    return response.choices[0].message.content.strip()

AGENTS = {
    "FAQ": {
        "system_prompt": "회사에 대한 자주 묻는 질문에 답변합니다.",
        "knowledge_base": "faq_docs/"
    },
    "TECHNICAL": {
        "system_prompt": "제품에 대한 기술 지원을 제공합니다.",
        "knowledge_base": "tech_docs/"
    },
    "POLICY": {
        "system_prompt": "회사 정책을 설명합니다. 절대 약속하지 마세요.",
        "knowledge_base": "policy_docs/"
    }
}

from chromadb import Client
 
def retrieve_context(query: str, collection_name: str, k: int = 3) -> str:
    collection = chroma_client.get_collection(collection_name)
    results = collection.query(query_texts=[query], n_results=k)
    return "\n".join(results["documents"][0])

def apply_guardrails(response: str, agent_type: str) -> tuple[bool, str]:
    # 정책 에이전트는 특정 결과를 약속해서는 안 됨
    if agent_type == "POLICY":
        forbidden = ["보장", "약속", "반드시", "100%"]
        for word in forbidden:
            if word.lower() in response.lower():
                return False, f"정책 위반: '{word}'"
 
    # PII 검사
    if contains_pii(response):
        return False, "응답에 PII 포함"
 
    return True, "OK"

class EvalResult(BaseModel):
    relevance: int = Field(ge=1, le=5)
    accuracy: int = Field(ge=1, le=5)
    safety: int = Field(ge=1, le=5)
    reasoning: str
 
def evaluate_response(query: str, response: str, context: str) -> EvalResult:
    completion = client.beta.chat.completions.parse(
        model="gpt-4o-mini",
        messages=[
            {"role": "system", "content": EVAL_PROMPT},
            {"role": "user", "content": f"쿼리: {query}\n컨텍스트: {context}\n응답: {response}"}
        ],
        response_format=EvalResult
    )
    return completion.choices[0].message.parsed

@app.post("/v1/knowledge/query")
async def knowledge_query(request: QueryRequest):
    # 1. 라우팅
    agent_type = route_query(request.query)
 
    # 2. 검색
    context = retrieve_context(request.query, AGENTS[agent_type]["knowledge_base"])
 
    # 3. 생성
    response = generate_response(request.query, context, agent_type)
 
    # 4. 가드레일
    is_safe, reason = apply_guardrails(response, agent_type)
    if not is_safe:
        return {"response": FALLBACK_RESPONSES[agent_type], "blocked": True}
 
    # 5. 평가 (비동기, 로깅용)
    asyncio.create_task(log_evaluation(request.query, response, context))
 
    return {"response": response, "agent": agent_type, "blocked": False}