src/rag/generator.py

"""
Generator module for Norwegian RAG chatbot.
Generates responses using retrieved context and LLM.
"""

from typing import List, Dict, Any, Optional

from ..api.huggingface_api import HuggingFaceAPI, create_rag_prompt

class Generator:
    """
    Generates responses using retrieved context and LLM.
    Uses Hugging Face Inference API for text generation.
    """
    
    def __init__(
        self,
        api_client: Optional[HuggingFaceAPI] = None,
    ):
        """
        Initialize the generator.
        
        Args:
            api_client: HuggingFaceAPI client for text generation
        """
        self.api_client = api_client or HuggingFaceAPI()
    
    def generate(
        self,
        query: str,
        retrieved_chunks: List[Dict[str, Any]],
        temperature: float = 0.7
    ) -> str:
        """
        Generate a response using retrieved context.
        
        Args:
            query: User query
            retrieved_chunks: List of retrieved chunks with metadata
            temperature: Temperature for text generation
            
        Returns:
            Generated response
        """
        # Extract text from retrieved chunks
        context_texts = [chunk["chunk_text"] for chunk in retrieved_chunks]
        
        # If no context is retrieved, generate a response without context
        if not context_texts:
            return self._generate_without_context(query, temperature)
        
        # Create RAG prompt
        prompt = create_rag_prompt(query, context_texts)
        
        # Generate response
        response = self.api_client.generate_text(
            prompt=prompt,
            temperature=temperature
        )
        
        return response
    
    def _generate_without_context(self, query: str, temperature: float = 0.7) -> str:
        """
        Generate a response without context when no relevant chunks are found.
        
        Args:
            query: User query
            temperature: Temperature for text generation
            
        Returns:
            Generated response
        """
        prompt = f"""Du er en hjelpsom assistent som svarer på norsk. Svar på følgende spørsmål så godt du kan.

SPØRSMÅL:
{query}

SVAR:
"""
        
        response = self.api_client.generate_text(
            prompt=prompt,
            temperature=temperature
        )
        
        return response