services/prompt_builder.py

"""

Prompt Builder Service for RAG Agent



Combines retrieved documentation chunks with user queries into well-formatted

prompts for OpenAI. Handles token management, XML formatting, and source attribution.



This service:

1. Formats context with XML delimiters

2. Manages token budgets to stay within API limits

3. Selects top-scoring chunks that fit

4. Includes source metadata for attribution

5. Adds system instructions for grounded answers



Usage:

    from services.prompt_builder import PromptBuilder



    builder = PromptBuilder(max_tokens=3000)

    prompt = builder.build_prompt(query="What is ROS 2?", chunks=results)

"""

import logging
from typing import List
from retrieval import SearchResult
from utils.token_counter import count_tokens

logger = logging.getLogger('prompt_builder')


class PromptBuilder:
    """

    Service for building prompts with context and query.



    This service creates well-formatted prompts that:

    - Use XML tags (<context>, <query>) for clear structure

    - Stay within token limits

    - Include source URLs for attribution

    - Provide system instructions for grounded answers



    Attributes:

        max_tokens: Maximum tokens for entire prompt (default 3000)

        reserved_tokens: Tokens reserved for query + instructions (default 500)



    Methods:

        build_prompt: Build formatted prompt from query and chunks

    """

    def __init__(self, max_tokens: int = 3000, reserved_tokens: int = 500):
        """

        Initialize prompt builder.



        Args:

            max_tokens: Maximum tokens for entire prompt (default 3000)

            reserved_tokens: Tokens reserved for query/instructions (default 500)

        """
        self.max_tokens = max_tokens
        self.reserved_tokens = reserved_tokens
        self.context_budget = max_tokens - reserved_tokens

        logger.info(
            f"PromptBuilder initialized",
            extra={
                "max_tokens": max_tokens,
                "reserved_tokens": reserved_tokens,
                "context_budget": self.context_budget
            }
        )

    def build_prompt(self, query: str, chunks: List[SearchResult], model: str = "gpt-3.5-turbo") -> str:
        """

        Build formatted prompt with context and query.



        Creates a prompt with the following structure:

        ```

        <context>

        [chunk 1 text]

        Source: [url 1]



        [chunk 2 text]

        Source: [url 2]

        ...

        </context>



        <query>

        [user query]

        </query>



        Please answer the query using ONLY the information provided...

        ```



        The method:

        1. Selects chunks that fit within token budget

        2. Formats each chunk with source URL

        3. Wraps in XML tags

        4. Adds system instructions



        Args:

            query: User's natural language question

            chunks: Retrieved documentation chunks (sorted by relevance)

            model: OpenAI model for token counting (default: gpt-3.5-turbo)



        Returns:

            Formatted prompt string ready for OpenAI API



        Example:

            >>> builder = PromptBuilder(max_tokens=3000)

            >>> prompt = builder.build_prompt("What is ROS 2?", chunks)

            >>> print(prompt[:100])

            <context>

            ROS 2 is a set of software libraries and tools...

        """
        logger.info(
            f"Building prompt",
            extra={
                "query": query[:50] + "..." if len(query) > 50 else query,
                "total_chunks": len(chunks),
                "model": model
            }
        )

        # Select chunks that fit within budget (T036)
        selected_chunks = self._select_chunks_within_budget(chunks, model)

        # Format context with chunks and source URLs (T037)
        context = self._format_context_with_sources(selected_chunks)

        # Build full prompt with XML template (T034)
        prompt = self._format_prompt_with_xml(context, query)

        # Count final tokens
        final_tokens = count_tokens(prompt, model)

        logger.info(
            f"Prompt built successfully",
            extra={
                "query": query[:50] + "..." if len(query) > 50 else query,
                "selected_chunks": len(selected_chunks),
                "total_chunks": len(chunks),
                "final_tokens": final_tokens,
                "max_tokens": self.max_tokens
            }
        )

        return prompt

    def _select_chunks_within_budget(

        self,

        chunks: List[SearchResult],

        model: str

    ) -> List[SearchResult]:
        """

        Select top-scoring chunks that fit within token budget.



        Iterates through chunks (already sorted by relevance) and adds

        them to context until budget is exhausted. Prioritizes higher-scoring

        chunks when space is limited.



        Args:

            chunks: Retrieved chunks sorted by relevance (highest first)

            model: Model for token counting



        Returns:

            List of chunks that fit within context_budget

        """
        selected = []
        total_tokens = 0

        for chunk in chunks:
            # Format chunk with source for accurate token count
            chunk_text = chunk.text
            source_url = chunk.metadata.get('url', '')
            formatted_chunk = f"{chunk_text}\nSource: {source_url}\n"

            # Count tokens for this chunk
            chunk_tokens = count_tokens(formatted_chunk, model)

            # Check if adding this chunk would exceed budget
            if total_tokens + chunk_tokens > self.context_budget:
                logger.warning(
                    f"Hit token budget, truncating context at {len(selected)} chunks",
                    extra={
                        "selected_chunks": len(selected),
                        "total_chunks": len(chunks),
                        "current_tokens": total_tokens,
                        "context_budget": self.context_budget
                    }
                )
                break

            selected.append(chunk)
            total_tokens += chunk_tokens

        logger.debug(
            f"Selected {len(selected)} chunks using {total_tokens} tokens"
        )

        return selected

    def _format_context_with_sources(self, chunks: List[SearchResult]) -> str:
        """

        Format chunks with source URLs.



        Each chunk is formatted as:

        ```

        [chunk text]

        Source: [url]

        ```



        Args:

            chunks: Selected chunks to format



        Returns:

            Formatted context string with sources

        """
        context_parts = []

        for chunk in chunks:
            chunk_text = chunk.text
            source_url = chunk.metadata.get('url', 'Unknown')

            # Format: chunk text + source URL
            formatted_chunk = f"{chunk_text}\nSource: {source_url}"
            context_parts.append(formatted_chunk)

        # Join with blank line between chunks
        context = "\n\n".join(context_parts)

        return context

    def _format_prompt_with_xml(self, context: str, query: str) -> str:
        """

        Format prompt with XML tags and system instructions.



        Creates the final prompt structure:

        ```

        <context>

        [formatted context with sources]

        </context>



        <query>

        [user query]

        </query>



        [system instructions]

        ```



        Args:

            context: Formatted context with chunks and sources

            query: User query



        Returns:

            Complete formatted prompt

        """
        # Build prompt with XML delimiters (T034)
        prompt = f"""<context>

{context}

</context>



<query>

{query}

</query>



Please answer the query using ONLY the information provided in the context above. If the context doesn't contain enough information to answer the query, say so. Be concise and cite specific details from the context."""

        return prompt