Merge pull request #51 from Team-StackUp/feature/ai-core-contract-question-rag

feat: 초기 질문/꼬리질문 생성에 Core embedding search 기반 RAG 추가

Author: pswaao88

ai/src/ai_server/chain/followup_generation_chain.py

@@ -27,6 +27,7 @@ async def generate(
         mode: str,
         previous_question: str,
         answer_text: str,
+        context: str = "(none)",
     ) -> FollowupResult: ...
 
 
@@ -41,13 +42,15 @@ async def generate(
         mode: str,
         previous_question: str,
         answer_text: str,
+        context: str = "(none)",
     ) -> FollowupResult:
         result = await self._chain.ainvoke(
             {
                 "job_category": job_category,
                 "mode": mode,
                 "previous_question": previous_question,
                 "answer_text": answer_text,
+                "context": context,
             }
         )
         if not isinstance(result, FollowupResult):

ai/src/ai_server/chain/prompts/followup_generation.py

@@ -18,7 +18,12 @@
 HUMAN_PROMPT = (
     "직군: {job_category}\n"
     "면접 모드: {mode}\n\n"
+    "모드별 지침:\n"
+    "- TECHNICAL: 기술 역량, 프로젝트 경험, 문제 해결을 중심으로 파고듭니다.\n"
+    "- PERSONALITY: 협업, 갈등 해결, 성장 경험을 중심으로 파고듭니다.\n"
+    "- INTEGRATED: 기술 질문과 인성 질문의 관점을 균형 있게 반영합니다.\n\n"
     "직전 질문:\n{previous_question}\n\n"
     "지원자 답변:\n{answer_text}\n\n"
+    "검색 문서 컨텍스트:\n---\n{context}\n---\n\n"
     "{format_instructions}"
 )

ai/src/ai_server/messaging/consumers/followup_consumer.py

@@ -4,13 +4,15 @@
 from aio_pika.abc import AbstractIncomingMessage
 
 from ai_server.chain.followup_generation_chain import FollowupGenerator
+from ai_server.core.client import CoreClient
 from ai_server.messaging.idempotency import LruIdempotencyStore
 from ai_server.messaging.publisher import CallbackPublisher
 from ai_server.model.envelope import Envelope
 from ai_server.model.messages.followup import (
     FollowupCallbackPayload,
     GenerateFollowupRequest,
 )
+from ai_server.rag.embedder import EmbeddingProvider
 
 log = structlog.get_logger(__name__)
 
@@ -23,11 +25,17 @@ def __init__(
         publisher: CallbackPublisher,
         idempotency: LruIdempotencyStore,
         callback_routing_key: str,
+        core_client: CoreClient | None = None,
+        embedder: EmbeddingProvider | None = None,
+        rag_top_k: int = 5,
     ) -> None:
         self._generator = generator
         self._publisher = publisher
         self._idempotency = idempotency
         self._callback_routing_key = callback_routing_key
+        self._core = core_client
+        self._embedder = embedder
+        self._rag_top_k = rag_top_k
 
     async def handle(self, message: AbstractIncomingMessage) -> None:
         async with message.process(requeue=False):
@@ -65,6 +73,7 @@ async def handle(self, message: AbstractIncomingMessage) -> None:
                 mode=req.mode,
                 previous_question=req.previous_question,
                 answer_text=req.answer_text,
+                context=await self._build_rag_context(req),
             )
 
             payload = FollowupCallbackPayload(
@@ -89,3 +98,25 @@ async def handle(self, message: AbstractIncomingMessage) -> None:
                 session_id=req.session_id,
                 trace_id=envelope.trace_id,
             )
+
+    async def _build_rag_context(self, req: GenerateFollowupRequest) -> str:
+        if not self._core or not self._embedder or not req.context_document_ids:
+            return "(none)"
+
+        query = f"{req.previous_question}\n\n{req.answer_text}"
+        try:
+            query_vec = (await self._embedder.embed([query]))[0]
+            hits = await self._core.search_embeddings(
+                query_embedding=query_vec,
+                document_ids=req.context_document_ids,
+                top_k=self._rag_top_k,
+            )
+        except Exception as exc:
+            log.warn("followup.rag.failed", error=str(exc), session_id=req.session_id)
+            return "(none)"
+
+        if not hits:
+            return "(none)"
+        return "\n---\n".join(
+            f"[doc#{h.document_id} chunk#{h.chunk_index}] {h.chunk_text}" for h in hits
+        )

ai/src/ai_server/messaging/consumers/questions_consumer.py

@@ -4,6 +4,7 @@
 from aio_pika.abc import AbstractIncomingMessage
 
 from ai_server.chain.question_generation_chain import QuestionGenerator
+from ai_server.core.client import CoreClient
 from ai_server.messaging.idempotency import LruIdempotencyStore
 from ai_server.messaging.publisher import CallbackPublisher
 from ai_server.model.envelope import Envelope
@@ -12,6 +13,7 @@
     GenerateQuestionsRequest,
     QuestionPoolCallbackPayload,
 )
+from ai_server.rag.embedder import EmbeddingProvider
 
 log = structlog.get_logger(__name__)
 
@@ -25,14 +27,20 @@ def __init__(
         idempotency: LruIdempotencyStore,
         callback_routing_key: str,
         initial_pool_size: int = 1,
+        core_client: CoreClient | None = None,
+        embedder: EmbeddingProvider | None = None,
+        rag_top_k: int = 5,
     ) -> None:
         self._generator = generator
         self._publisher = publisher
         self._idempotency = idempotency
         self._callback_routing_key = callback_routing_key
-        # Core 의 QuestionsCallbackService.applyPool 은 questions[0] 만 INSERT 하고 나머지는 폐기.
-        # 토큰 낭비를 줄이기 위해 envelope.max_questions 대신 풀 크기를 강제. 후속 작업에서 풀 저장 도입 시 늘리기 쉬움.
+        # Core compatibility keeps callback.kind=POOL, but this is the initial
+        # question result. maxQuestions remains the full session limit.
         self._initial_pool_size = max(1, initial_pool_size)
+        self._core = core_client
+        self._embedder = embedder
+        self._rag_top_k = rag_top_k
 
     async def handle(self, message: AbstractIncomingMessage) -> None:
         async with message.process(requeue=False):
@@ -57,7 +65,10 @@ async def handle(self, message: AbstractIncomingMessage) -> None:
                 return
 
             req = envelope.payload
-            effective_pool_size = self._initial_pool_size  # envelope.max_questions 무시
+            effective_pool_size = max(
+                1,
+                req.initial_question_count,
+            )
             log.info(
                 "questions.generate.start",
                 message_id=envelope.message_id,
@@ -68,7 +79,7 @@ async def handle(self, message: AbstractIncomingMessage) -> None:
                 trace_id=envelope.trace_id,
             )
 
-            context_text = _build_context(req.documents)
+            context_text = await self._build_context(req)
             pool = await self._generator.generate(
                 job_category=req.job_category,
                 mode=req.mode,
@@ -98,6 +109,34 @@ async def handle(self, message: AbstractIncomingMessage) -> None:
                 trace_id=envelope.trace_id,
             )
 
+    async def _build_context(self, req: GenerateQuestionsRequest) -> str:
+        base_context = _build_context(req.documents)
+        if not self._core or not self._embedder:
+            return base_context
+
+        document_ids = [d.document_id for d in req.documents]
+        if not document_ids:
+            return base_context
+
+        query = _build_initial_rag_query(req)
+        try:
+            query_vec = (await self._embedder.embed([query]))[0]
+            hits = await self._core.search_embeddings(
+                query_embedding=query_vec,
+                document_ids=document_ids,
+                top_k=self._rag_top_k,
+            )
+        except Exception as exc:
+            log.warn("questions.rag.failed", error=str(exc), session_id=req.session_id)
+            return base_context
+
+        if not hits:
+            return base_context
+        rag_context = "\n---\n".join(
+            f"[doc#{h.document_id} chunk#{h.chunk_index}] {h.chunk_text}" for h in hits
+        )
+        return f"{base_context}\n\n## Retrieved document chunks\n{rag_context}"
+
 
 def _build_context(documents: list[DocumentContext]) -> str:
     parts: list[str] = []
@@ -112,3 +151,20 @@ def _build_context(documents: list[DocumentContext]) -> str:
             block.append(d.markdown)
         parts.append("\n".join(block))
     return "\n\n".join(parts) if parts else "(no documents)"
+
+
+def _build_initial_rag_query(req: GenerateQuestionsRequest) -> str:
+    parts = [
+        f"mode: {req.mode}",
+        f"job category: {req.job_category}",
+    ]
+    for d in req.documents:
+        doc_parts = [f"document #{d.document_id} {d.source_type}"]
+        if d.summary:
+            doc_parts.append(d.summary)
+        if d.tech_stack:
+            doc_parts.append("tech stack: " + ", ".join(d.tech_stack))
+        if d.markdown:
+            doc_parts.append(d.markdown[:1000])
+        parts.append("\n".join(doc_parts))
+    return "\n\n".join(parts)

ai/src/ai_server/messaging/runner.py

@@ -154,6 +154,8 @@ def __init__(self, settings: Settings) -> None:
             idempotency=self._idempotency,
             callback_routing_key=settings.ai_callback_routing_questions,
             initial_pool_size=settings.questions_initial_pool_size,
+            core_client=core_client,
+            embedder=embedder,
         )
 
         # 꼬리질문 생성 (US-19)
@@ -165,6 +167,8 @@ def __init__(self, settings: Settings) -> None:
             publisher=self._publisher,
             idempotency=self._idempotency,
             callback_routing_key=settings.ai_callback_routing_questions,
+            core_client=core_client,
+            embedder=embedder,
         )
 
         # 종합 피드백 생성 (US-24)

ai/src/ai_server/model/messages/followup.py

@@ -1,6 +1,6 @@
 from typing import Literal
 
-from pydantic import BaseModel
+from pydantic import BaseModel, Field
 
 from ai_server.model._config import camel_config
 
@@ -18,6 +18,7 @@ class GenerateFollowupRequest(BaseModel):
     answer_text: str
     mode: InterviewMode
     job_category: Literal["FRONTEND", "BACKEND", "INFRA", "DBA"]
+    context_document_ids: list[int] = Field(default_factory=list)
 
 
 class AnswerEvaluation(BaseModel):

ai/src/ai_server/model/messages/questions.py

@@ -33,6 +33,7 @@ class GenerateQuestionsRequest(BaseModel):
     mode: InterviewMode
     job_category: JobCategory
     documents: list[DocumentContext] = []
+    initial_question_count: int = 1
     max_questions: int = 10
 
 

ai/tests/test_core_client.py

@@ -7,6 +7,7 @@
     CoreEmbeddingUpsertError,
     CoreTokenError,
     EmbeddingChunkPayload,
+    EmbeddingSearchHit,
     HttpCoreClient,
 )
 
@@ -241,3 +242,86 @@ async def test_upsert_embeddings_httpx_error_retriable() -> None:
         )
     assert exc_info.value.code == "CORE_UNAVAILABLE"
     assert exc_info.value.retriable is True
+
+
+# ----------- search_embeddings -----------
+
+
+def _make_post_client(
+    *,
+    status: int = 200,
+    json_body: dict | None = None,
+    text: str = "",
+    raise_exc: Exception | None = None,
+) -> MagicMock:
+    client = MagicMock()
+    resp = MagicMock(spec=httpx.Response)
+    resp.status_code = status
+    resp.text = text
+    resp.json = (
+        MagicMock(return_value=json_body)
+        if json_body is not None
+        else MagicMock(side_effect=ValueError("no json"))
+    )
+    if raise_exc is not None:
+        client.post = AsyncMock(side_effect=raise_exc)
+    else:
+        client.post = AsyncMock(return_value=resp)
+    return client
+
+
+@pytest.mark.asyncio
+async def test_search_embeddings_uses_latest_core_contract() -> None:
+    client = _make_post_client(
+        json_body={
+            "hits": [
+                {
+                    "documentId": 7,
+                    "chunkIndex": 2,
+                    "chunkText": "Spring transaction boundaries",
+                    "distance": 0.17,
+                }
+            ]
+        }
+    )
+    core = HttpCoreClient(base_url="http://core:38010", api_key="k", client=client)
+
+    hits = await core.search_embeddings(
+        query_embedding=[0.1, 0.2],
+        document_ids=[7, 8],
+        top_k=3,
+    )
+
+    assert hits == [
+        EmbeddingSearchHit(
+            document_id=7,
+            chunk_index=2,
+            chunk_text="Spring transaction boundaries",
+            distance=0.17,
+        )
+    ]
+    client.post.assert_awaited_once_with(
+        "/api/internal/embeddings/search",
+        json={
+            "queryEmbedding": [0.1, 0.2],
+            "documentIds": [7, 8],
+            "topK": 3,
+        },
+    )
+
+
+@pytest.mark.parametrize("status", [400, 401, 403, 404, 500])
+@pytest.mark.asyncio
+async def test_search_embeddings_non_2xx_returns_empty(status: int) -> None:
+    client = _make_post_client(status=status, text="bad")
+    core = HttpCoreClient(base_url="http://core:38010", api_key="k", client=client)
+
+    assert await core.search_embeddings(query_embedding=[0.1]) == []
+
+
+@pytest.mark.asyncio
+async def test_search_embeddings_http_error_returns_empty() -> None:
+    client = _make_post_client(raise_exc=httpx.ConnectError("dns fail"))
+    core = HttpCoreClient(base_url="http://core:38010", api_key="k", client=client)
+
+    assert await core.search_embeddings(query_embedding=[0.1]) == []