refactor(generation): fix imports, function names, and update comments for prompt building (#19)

Author: kernelshard

app/core/config.py

@@ -1,6 +1,8 @@
 from pydantic import Field
 from pydantic_settings import BaseSettings
 
+from app.core.constants import GeneratorBackend
+
 
 class Settings(BaseSettings):
     APP_NAME: str = "rag_mastery"
@@ -9,12 +11,15 @@ class Settings(BaseSettings):
     LOG_LEVEL: str = Field(default="INFO", alias="LOG_LEVEL")
     HOST: str = Field(default="0.0.0.0", alias="HOST")
     PORT: int = Field(default=8000, alias="PORT")
+    USE_REAL_GENERATOR: bool = False
+    GENERATOR_BACKEND: GeneratorBackend = GeneratorBackend.Mock
 
     model_config = {
         "env_file": ".env",
         "env_file_encoding": "utf-8",
         "extra": "ignore",
         "populate_by_name": True,  # allowing alias mapping
+        "use_enum_values": True,  # store enum values directly
     }
 
 

app/core/constants.py

@@ -1,6 +1,7 @@
 from enum import StrEnum
 
 
+# Ingestion status constants used across the application
 class IngestionStatus(StrEnum):
     Accepted = "accepted"
     Processing = "processing"
@@ -9,3 +10,14 @@ class IngestionStatus(StrEnum):
 
 
 DefaultTopK: int = 5
+
+
+class GeneratorBackend(StrEnum):
+    """
+    Supported generator backends.
+    Reason: to avoid hardcoding strings across the codebase.
+    """
+
+    Mock = "mock"
+    OPENAI = "openai"
+    Ollama = "ollama"

app/generation/adapters/ollama_adapter.py

@@ -0,0 +1,21 @@
+import asyncio
+from app.core.interfaces import BaseGenerator
+from app.core.logging import get_logger
+
+
+logger = get_logger(__name__)
+
+
+class OllamaGenerator(BaseGenerator):
+    """
+    Async adapter for Ollama-style LLMs.
+    """
+
+    def __init__(self, model_name: str = "llama3"):
+        self.model_name = model_name
+
+    async def generate(self, prompt: str) -> str:
+        logger.debug(f"[OllamaGenerator] Would call model={self.model_name}")
+        await asyncio.sleep(0.5)  # simulate network delay
+        # Placeholder for actual Ollama API call
+        return f"[Simulated Ollama: {self.model_name} for prompt: {prompt[:60]}...]"

app/generation/adapters/openai_adapter.py

@@ -0,0 +1,21 @@
+import asyncio
+
+from app.core.interfaces import BaseGenerator
+from app.core.logging import get_logger
+
+
+logger = get_logger(__name__)
+
+
+class OpenAIGenerator(BaseGenerator):
+    """
+    Async adapter for OpenAI-style LLMs.
+    """
+
+    def __init__(self, model_name: str = "gpt-4-turbo"):
+        self.model_name = model_name
+
+    async def generate(self, prompt: str) -> str:
+        logger.debug(f"[OpenAIGenerator] would call model={self.model_name}")
+        await asyncio.sleep(0.5)  # simulate network delay
+        return f"[Simulated OpenAI: {self.model_name} for prompt: {prompt[:60]}...]"

app/generation/api.py

@@ -4,7 +4,7 @@
 from app.core.logging import get_logger
 from app.generation.deps import get_retriever
 from app.generation.models import GenerationRequest, GenerationResponse
-from app.generation.service import generate_anwer
+from app.generation.service import generate_answer
 
 logger = get_logger(__name__)
 
@@ -23,4 +23,4 @@ async def generate_endpoint(
     logger.info(
         "Received generation request", query=req.query, context_size=req.context_size
     )
-    return await generate_anwer(req, retriever)
+    return await generate_answer(req, retriever)

app/generation/deps.py

@@ -1,8 +1,11 @@
+from app.core.constants import GeneratorBackend
 from app.core.interfaces import BaseRetriever
+from app.generation.adapters.ollama_adapter import OllamaGenerator
 from app.generation.mock_generator import MockGenerator
 from app.retrieval.models import RetrievalRequest
 from app.retrieval.service import retrieve_documents
 from app.core.repositories import global_vector_repo
+from app.core.config import settings
 
 
 class RetrievalAdapter(BaseRetriever):
@@ -35,6 +38,20 @@ def get_retriever() -> BaseRetriever:
     return _retriever
 
 
-async def get_generator():
-    # swap with real generator later
-    yield MockGenerator()
+async def get_generator(
+    use_real: bool = settings.USE_REAL_GENERATOR,
+    backend: GeneratorBackend = settings.GENERATOR_BACKEND,
+):
+    """
+    Returns a generator instance based on configuration.
+    """
+    # precedence wise use_real overrides backend
+    if not use_real:
+        yield MockGenerator()
+    else:
+        if backend == GeneratorBackend.Mock:
+            yield MockGenerator()
+        elif backend == GeneratorBackend.Ollama:
+            yield OllamaGenerator()
+        else:
+            yield OllamaGenerator()  # Default to OllamaGenerator for now

app/generation/prompt_builder.py

@@ -0,0 +1,29 @@
+def build_prompt(query: str, contexts: list[dict]) -> str:
+    """
+    Assembles a prompt text from retrieved contexts and the input query.
+
+    Expects contexts as list of dicts, e.g.:
+    [
+        {"doc_id": "1", "score": 0.8, "metadata": {"text": "FastAPI is async..."}},
+        ...
+    ]
+    """
+    if not contexts:
+        joined_contexts = "[No relevant context found.]"
+
+    else:
+        context_texts = []
+        for c in contexts:
+            meta = c.get("metadata", {})
+            text = meta.get("text")
+            if text:
+                context_texts.append(text.strip())
+            else:
+                context_texts.append(f"[Doc:{c.get('doc_id', 'unknown')}]")
+
+        joined_contexts = "\n\n".join(context_texts)
+    return (
+        f"### Contexts:\n{joined_contexts}\n\n"
+        f"### Question:\n{query.strip()}\n\n"
+        "### Answer:\n"
+    )

app/generation/service.py

@@ -1,15 +1,14 @@
-import asyncio
-
 from app.core.logging import get_logger
 from app.generation.models import GenerateAnswer, GenerationRequest, GenerationResponse
+from app.generation.prompt_builder import build_prompt
 
 from app.core.interfaces import BaseRetriever
 
 
 logger = get_logger(__name__)
 
 
-async def generate_anwer(
+async def generate_answer(
     req: GenerationRequest, retriever: BaseRetriever
 ) -> GenerationResponse:
     """
@@ -24,22 +23,21 @@ async def generate_anwer(
         GenerationResponse: The response containing the original query and generated answer.
     """
     # retrieve top-k relevant documents based on the query
+    logger.info(f"Generation started for query='{req.query}'")
+
     retrieved_chunks = await retriever.retrieve(req.query, req.context_size)
     """
     e.g: retrieved_chunks = [
-        {"doc_id": "1", "content": "Document content 1", "metadata": {"source": "source1"}},
-        {"doc_id": "2", "content": "Document content 2", "metadata": {"source": "source2"}},
+        {"doc_id": "1", "score": 0.8, "metadata": {"source": "source1"}},
+        {"doc_id": "2", "score": 0.7, "metadata": {"source": "source2"}},
         ...
     ]
+    Note: Does not include original text; only doc_id and metadata.
     """
+    logger.debug(f"Retrieved {len(retrieved_chunks)} chunks")
 
-    await asyncio.sleep(0.1)  # simulate llm latency
+    _ = build_prompt(req.query, retrieved_chunks)
 
-    # e.g context = [
-    #     {"doc_id": "1", "content": "Document content 1", "metadata": {"source": "source1"}},
-    #     {"doc_id": "2", "content": "Document content 2", "metadata": {"source": "source2"}},
-    #     ...
-    # ]
     synthesized = (
         " ".join(chunk["doc_id"] for chunk in retrieved_chunks) or "No context found."
     )

tests/generation/test_deps.py

@@ -0,0 +1,34 @@
+from app.core.constants import GeneratorBackend
+from app.generation.adapters.ollama_adapter import OllamaGenerator
+from app.generation.deps import get_generator
+from app.generation.mock_generator import MockGenerator
+
+
+import pytest
+
+
+@pytest.mark.asyncio
+async def test_get_generator_returns_mock_when_use_real_false_but_backend_not_provided():
+    """
+    When use_real is False and no backend is provided, should return MockGenerator.
+    """
+    gen = await anext(get_generator(use_real=False))
+    assert isinstance(gen, MockGenerator)
+
+
+@pytest.mark.asyncio
+async def test_get_generator_returns_mock_when_use_real_true_and_backend_mock():
+    """
+    When use_real is True and backend is Mock, should return MockGenerator.
+    """
+    gen = await anext(get_generator(use_real=True, backend=GeneratorBackend.Mock))
+    assert isinstance(gen, MockGenerator)
+
+
+@pytest.mark.asyncio
+async def test_get_generator_returns_ollama_when_use_real_false_and_backend_ollama():
+    """
+    When use_real is True and backend is Ollama, should return OllamaGenerator.
+    """
+    gen = await anext(get_generator(use_real=True, backend=GeneratorBackend.Ollama))
+    assert isinstance(gen, OllamaGenerator)