""" 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}")