drakkenheim-rag/src/drakkenheim/llm.py

"""
LLM and Vector Operations module for Drakkenheim RAG Application.

Handles OpenAI API interactions, embeddings, vector operations, and document re-ranking.
"""

import os
import openai
import chromadb
from sentence_transformers import CrossEncoder

# Custom OpenAI embedding function to avoid ChromaDB compatibility issues
class CustomOpenAIEmbeddingFunction:
    def __init__(self, api_key: str, model_name: str = "text-embedding-3-small"):
        self.api_key = api_key
        self.model_name = model_name
        self.client = openai.OpenAI(api_key=api_key)
    
    def __call__(self, input):
        """Generate embeddings for input texts."""
        try:
            # Truncate very long texts to avoid token limits
            max_chars = 8000  # Conservative limit for text-embedding-3-small
            truncated_input = []
            for text in input:
                if len(text) > max_chars:
                    truncated_input.append(text[:max_chars])
                    print(f"Warning: Truncated text from {len(text)} to {max_chars} characters")
                else:
                    truncated_input.append(text)

            response = self.client.embeddings.create(
                model=self.model_name,
                input=truncated_input
            )
            return [embedding.embedding for embedding in response.data]
        except Exception as e:
            print(f"Error generating embeddings: {e}")
            return [[0.0] * 1536 for _ in input]  # Return zero vectors as fallback

def get_vector_collection():
    """Gets or creates a ChromaDB collection for vector storage.

    Creates a custom OpenAI embedding function using the text-embedding-3-small model and initializes
    a persistent ChromaDB client. Returns a collection that can be used to store and
    query document embeddings.

    Returns:
        chromadb.Collection: A ChromaDB collection configured with the custom OpenAI embedding
            function and cosine similarity space.
    """
    openai_ef = CustomOpenAIEmbeddingFunction(
        api_key=os.getenv("OPENAI_API_KEY"),
        model_name="text-embedding-3-small",
    )

    # Use Docker volume path if available, fallback to local path
    chroma_path = os.getenv("CHROMA_DB_PATH", "./demo-rag-chroma")
    chroma_client = chromadb.PersistentClient(path=chroma_path)
    return chroma_client.get_or_create_collection(
        name="rag_app",
        embedding_function=openai_ef,
        metadata={"hnsw:space": "cosine"},
    )

def add_to_vector_collection(all_splits, file_name: str):
    """Adds document splits to a vector collection for semantic search.
    
    Processes document splits in batches to avoid token limits and adds them to the
    ChromaDB collection with proper metadata and unique IDs.
    
    Args:
        all_splits: List of Document objects to add to the collection
        file_name: Name of the file being processed (used for ID generation)
    """
    collection = get_vector_collection()
    batch_size = 50  # Process documents in smaller batches

    for batch_start in range(0, len(all_splits), batch_size):
        batch_end = min(batch_start + batch_size, len(all_splits))
        batch_splits = all_splits[batch_start:batch_end]

        documents, metadatas, ids = [], [], []

        for idx, split in enumerate(batch_splits):
            documents.append(split.page_content)
            metadatas.append(split.metadata)
            ids.append(f"{file_name}_{batch_start + idx}")

        collection.upsert(
            documents=documents,
            metadatas=metadatas,
            ids=ids,
        )

def query_collection(prompt: str, n_results: int = 10):
    """Queries the vector collection for relevant documents.
    
    Searches the ChromaDB collection for documents similar to the given prompt
    and returns the most relevant results.
    
    Args:
        prompt: The search query string
        n_results: Number of results to return (default: 10)
        
    Returns:
        dict: Query results containing documents, metadatas, distances, and ids
    """
    collection = get_vector_collection()
    results = collection.query(query_texts=[prompt], n_results=n_results)
    return results

def call_llm(context: str, prompt: str):
    """Calls the language model with context and prompt to generate a response.
    
    Uses OpenAI GPT-4o-mini API to stream responses from a language model by providing context and a
    question prompt. The model uses a system prompt to format and ground its responses appropriately.

    Args:
        context: String containing the relevant context for answering the question
        prompt: String containing the user's question

    Yields:
        String chunks of the generated response as they become available from the model

    Raises:
        OpenAIError: If there are issues communicating with the OpenAI API
    """
    client = openai.OpenAI(api_key=os.getenv("OPENAI_API_KEY"))
    
    response = client.chat.completions.create(
        model="gpt-4o-mini",
        stream=True,
        messages=[
            {
                "role": "system",
                "content": f"""You are an AI assistant tasked with providing detailed answers based solely on the given context. Your goal is to analyze the information provided and formulate a comprehensive, well-structured response to the question.

context will be passed as "Context:"
user question will be passed as "Question:"

Instructions:
1. Read the context carefully and identify the most relevant information
2. Use only the information provided in the context to answer the question
3. If the context doesn't contain enough information to answer the question, say so
4. Structure your response in a clear and organized manner
5. Cite specific parts of the context when relevant
6. Be concise but comprehensive

Context: {context}""",
            },
            {"role": "user", "content": prompt},
        ],
    )

    for chunk in response:
        if chunk.choices[0].delta.content is not None:
            yield chunk.choices[0].delta.content

def re_rank_cross_encoders(documents: list[str], prompt: str) -> tuple[str, list[int]]:
    """Re-ranks documents using a cross-encoder model for more accurate relevance scoring.
    
    Uses the MS MARCO MiniLM cross-encoder model to re-rank the input documents based on
    their relevance to the given prompt. This helps improve the quality of retrieved
    documents for RAG applications.
    
    Args:
        documents: List of document texts to re-rank
        prompt: The query prompt to rank documents against
        
    Returns:
        tuple: A tuple containing:
            - The most relevant document text
            - List of document indices in order of relevance
            
    Raises:
        RuntimeError: If cross-encoder model fails to load or rank documents
    """
    if not documents:
        return "", []
    
    try:
        encoder_model = CrossEncoder("cross-encoder/ms-marco-MiniLM-L-6-v2")
        ranks = encoder_model.rank(prompt, documents, top_k=3)
        relevant_text_ids = []
        relevant_text = ""
        
        for rank in ranks:
            relevant_text_ids.append(rank["corpus_id"])
            if not relevant_text:
                relevant_text = documents[rank["corpus_id"]]
        
        return relevant_text, relevant_text_ids
    except Exception as e:
        print(f"Error in cross-encoder ranking: {e}")
        # Fallback to returning first document
        return documents[0] if documents else "", [0]