""" 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]