llm_handler.py

"""
LLM Handler Module
Manages interaction with Hugging Face Inference API for answer generation
"""
import os
from typing import Generator, Dict, List
import logging
from huggingface_hub import InferenceClient

from config import (
    HF_MODEL,
    HF_TOKEN,
    SYSTEM_PROMPT,
    PROMPT_TEMPLATE,
)

# Configure logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)


class LLMHandler:
    """
    Handles LLM interactions using Hugging Face Inference API
    """
    
    def __init__(self, model: str = HF_MODEL, token: str = None):
        """
        Initialize the LLM handler
        
        Args:
            model: Name of the Hugging Face model to use
            token: HF API token (if not provided, will use HF_TOKEN from config)
        """
        self.model = model
        self.token = token or HF_TOKEN
        
        if not self.token:
            raise ValueError(
                "Hugging Face token not found. Please set HF_TOKEN environment variable "
                "or pass it to the constructor."
            )
        
        self.client = InferenceClient(token=self.token)
        logger.info(f"Initialized LLM handler with model: {model}")
    
    def generate_answer(
        self,
        question: str,
        context: str,
        stream: bool = False
    ) -> str:
        """
        Generate an answer based on the question and context
        
        Args:
            question: User's question
            context: Retrieved context from documents
            stream: Whether to stream the response
            
        Returns:
            Generated answer
        """
        # Format the prompt with system prompt, context, and question
        full_prompt = f"{SYSTEM_PROMPT}\n\n{PROMPT_TEMPLATE.format(context=context, question=question)}"
        
        try:
            if not stream:
                response = self.client.text_generation(
                    prompt=full_prompt,
                    model=self.model,
                    max_new_tokens=1024,
                    temperature=0.7,
                    top_p=0.95,
                    stream=False,
                )
                return response
            else:
                # Return generator for streaming
                return self.client.text_generation(
                    prompt=full_prompt,
                    model=self.model,
                    max_new_tokens=1024,
                    temperature=0.7,
                    top_p=0.95,
                    stream=True,
                )
        
        except Exception as e:
            logger.error(f"Error generating answer: {e}")
            raise
    
    def stream_answer(
        self,
        question: str,
        context: str
    ) -> Generator[str, None, None]:
        """
        Stream the answer generation token by token
        
        Args:
            question: User's question
            context: Retrieved context from documents
            
        Yields:
            Generated text tokens
        """
        # Format the prompt with system prompt, context, and question
        full_prompt = f"{SYSTEM_PROMPT}\n\n{PROMPT_TEMPLATE.format(context=context, question=question)}"
        
        try:
            stream = self.client.text_generation(
                prompt=full_prompt,
                model=self.model,
                max_new_tokens=512,
                temperature=0.7,
                top_p=0.95,
                stream=True,
                details=False,
            )
            
            # Collect tokens and yield them
            has_content = False
            for token in stream:
                if token:  # Only yield non-empty tokens
                    has_content = True
                    yield token
            
            # If no content was generated, yield a fallback message
            if not has_content:
                yield "I apologize, but I couldn't generate an answer based on the provided context. Please try rephrasing your question."
        
        except StopIteration:
            # Handle empty generator
            yield "I apologize, but I couldn't generate an answer. The model returned an empty response."
        except Exception as e:
            logger.error(f"Error streaming answer: {e}")
            yield f"\n\n❌ Error: {str(e)}"


def format_response(
    question: str,
    answer: str,
    sources: List[Dict]
) -> str:
    """
    Format the final response with question, answer, and sources
    
    Args:
        question: User's question
        answer: Generated answer
        sources: List of source documents
        
    Returns:
        Formatted response in markdown
    """
    # Create response header
    response_parts = [
        f"**Question:** {question}\n",
        f"**Answer:** {answer}\n",
    ]
    
    # Add sources section
    if sources:
        response_parts.append("\n**Sources:**\n")
        
        # Group sources by document
        sources_by_doc = {}
        for source in sources:
            doc_name = source["source"]
            if doc_name not in sources_by_doc:
                sources_by_doc[doc_name] = []
            sources_by_doc[doc_name].append(source)
        
        # Format sources
        for doc_name, doc_sources in sources_by_doc.items():
            chunks = ", ".join([s["chunk_id"] for s in doc_sources])
            avg_similarity = sum(s["similarity"] for s in doc_sources) / len(doc_sources)
            response_parts.append(
                f"- {doc_name} (chunks: {chunks}, "
                f"relevance: {avg_similarity:.2%})\n"
            )
    
    return "".join(response_parts)


def stream_llm_answer(
    question: str,
    context: str
) -> Generator[str, None, None]:
    """
    Stream answer generation for a question with context
    
    Args:
        question: User's question
        context: Retrieved context
        
    Yields:
        Generated text tokens
    """
    llm = LLMHandler()
    
    try:
        for token in llm.stream_answer(question, context):
            yield token
    except Exception as e:
        logger.error(f"Error in stream_llm_answer: {e}")
        yield f"\n\n❌ Error generating answer: {str(e)}"


def generate_answer(
    question: str,
    context: str
) -> str:
    """
    Generate a complete answer for a question with context
    
    Args:
        question: User's question
        context: Retrieved context
        
    Returns:
        Generated answer
    """
    llm = LLMHandler()
    
    try:
        answer = llm.generate_answer(question, context, stream=False)
        return answer
    except Exception as e:
        logger.error(f"Error generating answer: {e}")
        return f"❌ Error generating answer: {str(e)}"


if __name__ == "__main__":
    # Test the LLM handler
    logger.info("Testing LLM handler...")
    
    # Create LLM instance
    llm = LLMHandler()
    
    # Test answer generation
    test_question = "What is Python?"
    test_context = "Python is a high-level programming language known for its simplicity and readability."
    
    logger.info(f"\nTest Question: {test_question}")
    logger.info(f"Context: {test_context}\n")
    
    # Test streaming
    logger.info("Streaming answer:")
    for token in llm.stream_answer(test_question, test_context):
        print(token, end='', flush=True)
    print("\n")
    
    # Test non-streaming
    logger.info("\nGenerating complete answer:")
    answer = llm.generate_answer(test_question, test_context)
    logger.info(f"Answer: {answer}")